import re import json import time import copy import elasticsearch from elastic_transport import ConnectionTimeout from elasticsearch import Elasticsearch from elasticsearch_dsl import UpdateByQuery, Search, Index from rag.settings import es_logger from rag import settings from rag.utils import singleton es_logger.info("Elasticsearch version: "+str(elasticsearch.__version__)) @singleton class ESConnection: def __init__(self): self.info = {} self.conn() self.idxnm = settings.ES.get("index_name", "") if not self.es.ping(): raise Exception("Can't connect to ES cluster") def conn(self): for _ in range(10): try: self.es = Elasticsearch( settings.ES["hosts"].split(","), basic_auth=(settings.ES["username"], settings.ES["password"]) if "username" in settings.ES and "password" in settings.ES else None, verify_certs=False, timeout=600 ) if self.es: self.info = self.es.info() es_logger.info("Connect to es.") break except Exception as e: es_logger.error("Fail to connect to es: " + str(e)) time.sleep(1) def version(self): v = self.info.get("version", {"number": "5.6"}) v = v["number"].split(".")[0] return int(v) >= 7 def health(self): return dict(self.es.cluster.health()) def upsert(self, df, idxnm=""): res = [] for d in df: id = d["id"] del d["id"] d = {"doc": d, "doc_as_upsert": "true"} T = False for _ in range(10): try: if not self.version(): r = self.es.update( index=( self.idxnm if not idxnm else idxnm), body=d, id=id, doc_type="doc", refresh=True, retry_on_conflict=100) else: r = self.es.update( index=( self.idxnm if not idxnm else idxnm), body=d, id=id, refresh=True, retry_on_conflict=100) es_logger.info("Successfully upsert: %s" % id) T = True break except Exception as e: es_logger.warning("Fail to index: " + json.dumps(d, ensure_ascii=False) + str(e)) if re.search(r"(Timeout|time out)", str(e), re.IGNORECASE): time.sleep(3) continue self.conn() T = False if not T: res.append(d) es_logger.error( "Fail to index: " + re.sub( "[\r\n]", "", json.dumps( d, ensure_ascii=False))) d["id"] = id d["_index"] = self.idxnm if not res: return True return False def bulk(self, df, idx_nm=None): ids, acts = {}, [] for d in df: id = d["id"] if "id" in d else d["_id"] ids[id] = copy.deepcopy(d) ids[id]["_index"] = self.idxnm if not idx_nm else idx_nm if "id" in d: del d["id"] if "_id" in d: del d["_id"] acts.append( {"update": {"_id": id, "_index": ids[id]["_index"]}, "retry_on_conflict": 100}) acts.append({"doc": d, "doc_as_upsert": "true"}) res = [] for _ in range(100): try: if elasticsearch.__version__[0] < 8: r = self.es.bulk( index=( self.idxnm if not idx_nm else idx_nm), body=acts, refresh=False, timeout="600s") else: r = self.es.bulk(index=(self.idxnm if not idx_nm else idx_nm), operations=acts, refresh=False, timeout="600s") if re.search(r"False", str(r["errors"]), re.IGNORECASE): return res for it in r["items"]: if "error" in it["update"]: res.append(str(it["update"]["_id"]) + ":" + str(it["update"]["error"])) return res except Exception as e: es_logger.warn("Fail to bulk: " + str(e)) if re.search(r"(Timeout|time out)", str(e), re.IGNORECASE): time.sleep(3) continue self.conn() return res def bulk4script(self, df): ids, acts = {}, [] for d in df: id = d["id"] ids[id] = copy.deepcopy(d["raw"]) acts.append({"update": {"_id": id, "_index": self.idxnm}}) acts.append(d["script"]) es_logger.info("bulk upsert: %s" % id) res = [] for _ in range(10): try: if not self.version(): r = self.es.bulk( index=self.idxnm, body=acts, refresh=False, timeout="600s", doc_type="doc") else: r = self.es.bulk( index=self.idxnm, body=acts, refresh=False, timeout="600s") if re.search(r"False", str(r["errors"]), re.IGNORECASE): return res for it in r["items"]: if "error" in it["update"]: res.append(str(it["update"]["_id"])) return res except Exception as e: es_logger.warning("Fail to bulk: " + str(e)) if re.search(r"(Timeout|time out)", str(e), re.IGNORECASE): time.sleep(3) continue self.conn() return res def rm(self, d): for _ in range(10): try: if not self.version(): r = self.es.delete( index=self.idxnm, id=d["id"], doc_type="doc", refresh=True) else: r = self.es.delete( index=self.idxnm, id=d["id"], refresh=True, doc_type="_doc") es_logger.info("Remove %s" % d["id"]) return True except Exception as e: es_logger.warn("Fail to delete: " + str(d) + str(e)) if re.search(r"(Timeout|time out)", str(e), re.IGNORECASE): time.sleep(3) continue if re.search(r"(not_found)", str(e), re.IGNORECASE): return True self.conn() es_logger.error("Fail to delete: " + str(d)) return False def search(self, q, idxnm=None, src=False, timeout="2s"): if not isinstance(q, dict): q = Search().query(q).to_dict() for i in range(3): try: res = self.es.search(index=(self.idxnm if not idxnm else idxnm), body=q, timeout=timeout, # search_type="dfs_query_then_fetch", track_total_hits=True, _source=src) if str(res.get("timed_out", "")).lower() == "true": raise Exception("Es Timeout.") return res except Exception as e: es_logger.error( "ES search exception: " + str(e) + "【Q】：" + str(q)) if str(e).find("Timeout") > 0: continue raise e es_logger.error("ES search timeout for 3 times!") raise Exception("ES search timeout.") def sql(self, sql, fetch_size=128, format="json", timeout="2s"): for i in range(3): try: res = self.es.sql.query(body={"query": sql, "fetch_size": fetch_size}, format=format, request_timeout=timeout) return res except ConnectionTimeout as e: es_logger.error("Timeout【Q】：" + sql) continue except Exception as e: raise e es_logger.error("ES search timeout for 3 times!") raise ConnectionTimeout() def get(self, doc_id, idxnm=None): for i in range(3): try: res = self.es.get(index=(self.idxnm if not idxnm else idxnm), id=doc_id) if str(res.get("timed_out", "")).lower() == "true": raise Exception("Es Timeout.") return res except Exception as e: es_logger.error( "ES get exception: " + str(e) + "【Q】：" + doc_id) if str(e).find("Timeout") > 0: continue raise e es_logger.error("ES search timeout for 3 times!") raise Exception("ES search timeout.") def updateByQuery(self, q, d): ubq = UpdateByQuery(index=self.idxnm).using(self.es).query(q) scripts = "" for k, v in d.items(): scripts += "ctx._source.%s = params.%s;" % (str(k), str(k)) ubq = ubq.script(source=scripts, params=d) ubq = ubq.params(refresh=False) ubq = ubq.params(slices=5) ubq = ubq.params(conflicts="proceed") for i in range(3): try: r = ubq.execute() return True except Exception as e: es_logger.error("ES updateByQuery exception: " + str(e) + "【Q】：" + str(q.to_dict())) if str(e).find("Timeout") > 0 or str(e).find("Conflict") > 0: continue self.conn() return False def updateScriptByQuery(self, q, scripts, idxnm=None): ubq = UpdateByQuery( index=self.idxnm if not idxnm else idxnm).using( self.es).query(q) ubq = ubq.script(source=scripts) ubq = ubq.params(refresh=True) ubq = ubq.params(slices=5) ubq = ubq.params(conflicts="proceed") for i in range(3): try: r = ubq.execute() return True except Exception as e: es_logger.error("ES updateByQuery exception: " + str(e) + "【Q】：" + str(q.to_dict())) if str(e).find("Timeout") > 0 or str(e).find("Conflict") > 0: continue self.conn() return False def deleteByQuery(self, query, idxnm=""): for i in range(3): try: r = self.es.delete_by_query( index=idxnm if idxnm else self.idxnm, refresh = True, body=Search().query(query).to_dict()) return True except Exception as e: es_logger.error("ES updateByQuery deleteByQuery: " + str(e) + "【Q】：" + str(query.to_dict())) if str(e).find("NotFoundError") > 0: return True if str(e).find("Timeout") > 0 or str(e).find("Conflict") > 0: continue return False def update(self, id, script, routing=None): for i in range(3): try: if not self.version(): r = self.es.update( index=self.idxnm, id=id, body=json.dumps( script, ensure_ascii=False), doc_type="doc", routing=routing, refresh=False) else: r = self.es.update(index=self.idxnm, id=id, body=json.dumps(script, ensure_ascii=False), routing=routing, refresh=False) # , doc_type="_doc") return True except Exception as e: es_logger.error( "ES update exception: " + str(e) + " id：" + str(id) + ", version:" + str(self.version()) + json.dumps(script, ensure_ascii=False)) if str(e).find("Timeout") > 0: continue return False def indexExist(self, idxnm): s = Index(idxnm if idxnm else self.idxnm, self.es) for i in range(3): try: return s.exists() except Exception as e: es_logger.error("ES updateByQuery indexExist: " + str(e)) if str(e).find("Timeout") > 0 or str(e).find("Conflict") > 0: continue return False def docExist(self, docid, idxnm=None): for i in range(3): try: return self.es.exists(index=(idxnm if idxnm else self.idxnm), id=docid) except Exception as e: es_logger.error("ES Doc Exist: " + str(e)) if str(e).find("Timeout") > 0 or str(e).find("Conflict") > 0: continue return False def createIdx(self, idxnm, mapping): try: if elasticsearch.__version__[0] < 8: return self.es.indices.create(idxnm, body=mapping) from elasticsearch.client import IndicesClient return IndicesClient(self.es).create(index=idxnm, settings=mapping["settings"], mappings=mapping["mappings"]) except Exception as e: es_logger.error("ES create index error %s ----%s" % (idxnm, str(e))) def deleteIdx(self, idxnm): try: return self.es.indices.delete(idxnm, allow_no_indices=True) except Exception as e: es_logger.error("ES delete index error %s ----%s" % (idxnm, str(e))) def getTotal(self, res): if isinstance(res["hits"]["total"], type({})): return res["hits"]["total"]["value"] return res["hits"]["total"] def getDocIds(self, res): return [d["_id"] for d in res["hits"]["hits"]] def getSource(self, res): rr = [] for d in res["hits"]["hits"]: d["_source"]["id"] = d["_id"] d["_source"]["_score"] = d["_score"] rr.append(d["_source"]) return rr def scrollIter(self, pagesize=100, scroll_time='2m', q={ "query": {"match_all": {}}, "sort": [{"updated_at": {"order": "desc"}}]}): for _ in range(100): try: page = self.es.search( index=self.idxnm, scroll=scroll_time, size=pagesize, body=q, _source=None ) break except Exception as e: es_logger.error("ES scrolling fail. " + str(e)) time.sleep(3) sid = page['_scroll_id'] scroll_size = page['hits']['total']["value"] es_logger.info("[TOTAL]%d" % scroll_size) # Start scrolling while scroll_size > 0: yield page["hits"]["hits"] for _ in range(100): try: page = self.es.scroll(scroll_id=sid, scroll=scroll_time) break except Exception as e: es_logger.error("ES scrolling fail. " + str(e)) time.sleep(3) # Update the scroll ID sid = page['_scroll_id'] # Get the number of results that we returned in the last scroll scroll_size = len(page['hits']['hits']) ELASTICSEARCH = ESConnection()