mirror of
https://github.com/tubearchivist/tubearchivist.git
synced 2025-02-18 14:00:21 +00:00
Refactor IndexPaginate
- add callback function - implement callback for ElasticBackup
This commit is contained in:
parent
3b17c01c6d
commit
6bb7f80ea2
@ -91,16 +91,22 @@ class ElasticWrap:
|
||||
|
||||
|
||||
class IndexPaginate:
|
||||
"""use search_after to go through whole index"""
|
||||
"""use search_after to go through whole index
|
||||
kwargs:
|
||||
- size: int, overwrite DEFAULT_SIZE
|
||||
- keep_source: bool, keep _source key from es resutls
|
||||
- callback: obj, Class with run method collback for every loop
|
||||
"""
|
||||
|
||||
DEFAULT_SIZE = 500
|
||||
|
||||
def __init__(self, index_name, data, size=False, keep_source=False):
|
||||
def __init__(self, index_name, data, **kwargs):
|
||||
self.index_name = index_name
|
||||
self.data = data
|
||||
self.pit_id = False
|
||||
self.size = size
|
||||
self.keep_source = keep_source
|
||||
self.size = kwargs.get("size")
|
||||
self.keep_source = kwargs.get("keep_source")
|
||||
self.callback = kwargs.get("callback")
|
||||
|
||||
def get_results(self):
|
||||
"""get all results"""
|
||||
@ -122,14 +128,13 @@ class IndexPaginate:
|
||||
print(self.data)
|
||||
raise ValueError("missing sort key in data")
|
||||
|
||||
size = self.size or self.DEFAULT_SIZE
|
||||
|
||||
self.data["size"] = size
|
||||
self.data["size"] = self.size or self.DEFAULT_SIZE
|
||||
self.data["pit"] = {"id": self.pit_id, "keep_alive": "10m"}
|
||||
|
||||
def run_loop(self):
|
||||
"""loop through results until last hit"""
|
||||
all_results = []
|
||||
counter = 0
|
||||
while True:
|
||||
response, _ = ElasticWrap("_search").get(data=self.data)
|
||||
all_hits = response["hits"]["hits"]
|
||||
@ -139,10 +144,18 @@ class IndexPaginate:
|
||||
source = hit
|
||||
else:
|
||||
source = hit["_source"]
|
||||
search_after = hit["sort"]
|
||||
all_results.append(source)
|
||||
|
||||
if not self.callback:
|
||||
all_results.append(source)
|
||||
|
||||
if self.callback:
|
||||
self.callback(all_hits, self.index_name).run()
|
||||
if counter % 10 == 0:
|
||||
print(f"{self.index_name}: processing page {counter}")
|
||||
counter = counter + 1
|
||||
|
||||
# update search_after with last hit data
|
||||
self.data["search_after"] = search_after
|
||||
self.data["search_after"] = all_hits[-1]["sort"]
|
||||
else:
|
||||
break
|
||||
|
||||
|
@ -137,35 +137,24 @@ class ElasticIndex:
|
||||
_, _ = ElasticWrap(f"ta_{self.index_name}").put(data)
|
||||
|
||||
|
||||
class ElasticBackup:
|
||||
"""dump index to nd-json files for later bulk import"""
|
||||
class BackupCallback:
|
||||
"""handle backup ndjson writer as callback for IndexPaginate"""
|
||||
|
||||
def __init__(self, index_config, reason):
|
||||
self.config = AppConfig().config
|
||||
self.cache_dir = self.config["application"]["cache_dir"]
|
||||
self.index_config = index_config
|
||||
self.reason = reason
|
||||
def __init__(self, source, index_name):
|
||||
self.source = source
|
||||
self.index_name = index_name
|
||||
self.timestamp = datetime.now().strftime("%Y%m%d")
|
||||
self.backup_files = []
|
||||
|
||||
@staticmethod
|
||||
def get_all_documents(index_name):
|
||||
"""export all documents of a single index"""
|
||||
data = {
|
||||
"query": {"match_all": {}},
|
||||
"sort": [{"_doc": {"order": "desc"}}],
|
||||
}
|
||||
paginate = IndexPaginate(f"ta_{index_name}", data, keep_source=True)
|
||||
all_results = paginate.get_results()
|
||||
def run(self):
|
||||
"""run the junk task"""
|
||||
file_content = self._build_bulk()
|
||||
self._write_es_json(file_content)
|
||||
|
||||
return all_results
|
||||
|
||||
@staticmethod
|
||||
def build_bulk(all_results):
|
||||
def _build_bulk(self):
|
||||
"""build bulk query data from all_results"""
|
||||
bulk_list = []
|
||||
|
||||
for document in all_results:
|
||||
for document in self.source:
|
||||
document_id = document["_id"]
|
||||
es_index = document["_index"]
|
||||
action = {"index": {"_index": es_index, "_id": document_id}}
|
||||
@ -179,29 +168,56 @@ class ElasticBackup:
|
||||
|
||||
return file_content
|
||||
|
||||
def write_es_json(self, file_content, index_name):
|
||||
def _write_es_json(self, file_content):
|
||||
"""write nd-json file for es _bulk API to disk"""
|
||||
file_name = f"es_{index_name}-{self.timestamp}.json"
|
||||
file_path = os.path.join(self.cache_dir, "backup", file_name)
|
||||
with open(file_path, "w", encoding="utf-8") as f:
|
||||
cache_dir = AppConfig().config["application"]["cache_dir"]
|
||||
file_name = f"es_{self.index_name.lstrip('ta_')}-{self.timestamp}.json"
|
||||
file_path = os.path.join(cache_dir, "backup", file_name)
|
||||
with open(file_path, "a+", encoding="utf-8") as f:
|
||||
f.write(file_content)
|
||||
|
||||
self.backup_files.append(file_path)
|
||||
|
||||
class ElasticBackup:
|
||||
"""dump index to nd-json files for later bulk import"""
|
||||
|
||||
def __init__(self, index_config, reason):
|
||||
self.config = AppConfig().config
|
||||
self.cache_dir = self.config["application"]["cache_dir"]
|
||||
self.timestamp = datetime.now().strftime("%Y%m%d")
|
||||
self.index_config = index_config
|
||||
self.reason = reason
|
||||
|
||||
@staticmethod
|
||||
def backup_index(index_name):
|
||||
"""export all documents of a single index"""
|
||||
data = {
|
||||
"query": {"match_all": {}},
|
||||
"sort": [{"_doc": {"order": "desc"}}],
|
||||
}
|
||||
paginate = IndexPaginate(
|
||||
f"ta_{index_name}", data, keep_source=True, callback=BackupCallback
|
||||
)
|
||||
_ = paginate.get_results()
|
||||
|
||||
def zip_it(self):
|
||||
"""pack it up into single zip file"""
|
||||
file_name = f"ta_backup-{self.timestamp}-{self.reason}.zip"
|
||||
backup_folder = os.path.join(self.cache_dir, "backup")
|
||||
backup_file = os.path.join(backup_folder, file_name)
|
||||
folder = os.path.join(self.cache_dir, "backup")
|
||||
|
||||
with zipfile.ZipFile(
|
||||
backup_file, "w", compression=zipfile.ZIP_DEFLATED
|
||||
) as zip_f:
|
||||
for backup_file in self.backup_files:
|
||||
to_backup = []
|
||||
for file in os.listdir(folder):
|
||||
if file.endswith(".json"):
|
||||
to_backup.append(os.path.join(folder, file))
|
||||
|
||||
backup_file = os.path.join(folder, file_name)
|
||||
|
||||
comp = zipfile.ZIP_DEFLATED
|
||||
with zipfile.ZipFile(backup_file, "w", compression=comp) as zip_f:
|
||||
for backup_file in to_backup:
|
||||
zip_f.write(backup_file, os.path.basename(backup_file))
|
||||
|
||||
# cleanup
|
||||
for backup_file in self.backup_files:
|
||||
for backup_file in to_backup:
|
||||
os.remove(backup_file)
|
||||
|
||||
def post_bulk_restore(self, file_name):
|
||||
@ -369,9 +385,7 @@ def backup_all_indexes(reason):
|
||||
print(f"backup: export in progress for {index_name}")
|
||||
if not backup_handler.index_exists(index_name):
|
||||
continue
|
||||
all_results = backup_handler.get_all_documents(index_name)
|
||||
file_content = backup_handler.build_bulk(all_results)
|
||||
backup_handler.write_es_json(file_content, index_name)
|
||||
backup_handler.backup_index(index_name)
|
||||
|
||||
backup_handler.zip_it()
|
||||
|
||||
|
Loading…
Reference in New Issue
Block a user