From 9779dfd4c5b7e834180f3257f9c64f0e77e473eb Mon Sep 17 00:00:00 2001 From: Michael Granitzer Date: Thu, 14 Nov 2024 16:36:22 +0100 Subject: [PATCH] feat: opensearch plug plus converter alpha version --- owilix/_version.py | 4 +- owilix/plugins/push/convert.py | 16 ++++ owilix/plugins/push/opensearch.py | 137 ++++++++++++++++++------------ pyproject.toml | 1 + 4 files changed, 100 insertions(+), 58 deletions(-) create mode 100644 owilix/plugins/push/convert.py diff --git a/owilix/_version.py b/owilix/_version.py index 8aba27f..6b4a99b 100644 --- a/owilix/_version.py +++ b/owilix/_version.py @@ -1,3 +1,3 @@ # These version placeholders will be replaced later during substitution. -__version__ = "0.14.0" -__version_tuple__ = (0, 14, 0) +__version__ = "0.14.0-post.1+004060a" +__version_tuple__ = (0, 14, 0, "post", 1, "004060a") diff --git a/owilix/plugins/push/convert.py b/owilix/plugins/push/convert.py new file mode 100644 index 0000000..95ee781 --- /dev/null +++ b/owilix/plugins/push/convert.py @@ -0,0 +1,16 @@ +from datetime import datetime +from typing import List, Dict +from dateutil import parser + +class DateConverter: + def __init__(self, **kwargs): + pass + + def __call__(self, jsons:List[Dict])->List[Dict]: + """ expect list of dicts and returns a list of dicts that must be JSON Serializable (i.e. no complex data types""" + for json in jsons: + if "warc_date" in json and isinstance(json["warc_date"],str): + json["warc_date"] =parser.parse(json["warc_date"]).strftime('%Y-%m-%dT%H:%M:%SZ') + if "day" in json and "month" in json and "year" in json: + json["date"] = datetime(json["year"],json["month"],json["day"]).strftime('%Y-%m-%dT%H:%M:%SZ') + return jsons \ No newline at end of file diff --git a/owilix/plugins/push/opensearch.py b/owilix/plugins/push/opensearch.py index 1302df9..6243839 100644 --- a/owilix/plugins/push/opensearch.py +++ b/owilix/plugins/push/opensearch.py @@ -1,3 +1,4 @@ +import importlib import json import os import sys @@ -108,7 +109,12 @@ class OpenSearchIndexer(BaseCommand): self.console.log(f"[red]Error communicating with OpenSearch: {e}[/red]") sys.exit(1) - def push_to_opensearch(self, index_name, batch_size=1000, delete_first=False, settings=None, mappings=None, file_path=None): + def push_to_opensearch(self, index_name, batch_size=1000, + delete_first=False, + settings=None, + mappings=None, + file_path=None, + fn_convert=None): """ Reads JSON data from stdin or a file and indexes it in batches to OpenSearch. @@ -150,61 +156,64 @@ class OpenSearchIndexer(BaseCommand): for line in data_source: try: - doc = json.loads(line.strip()) - # Allow per-document _op_type (e.g., 'index', 'create', 'update', 'delete') - op_type = doc.pop('_op_type', 'index') - action_metadata = {"_index": index_name} - if '_id' in doc: - action_metadata['_id'] = doc.pop('_id') - action_line = {op_type: action_metadata} - action_line_json = json.dumps(action_line) - if op_type == 'delete': - # For delete operation, no source needed - actions.append(action_line_json) - action_count += 1 - elif op_type == 'update': - # For update operation, wrap the doc in 'doc' - doc_line_json = json.dumps({'doc': doc}) - actions.append(action_line_json) - actions.append(doc_line_json) - action_count += 1 - else: - # For index or create operation - doc_line_json = json.dumps(doc) - actions.append(action_line_json) - actions.append(doc_line_json) - action_count += 1 - - progress.update(task, advance=1) - - # Send batch to OpenSearch when batch_size is reached - if action_count >= batch_size: - # Prepare bulk data - bulk_data = '\n'.join(actions) + '\n' # Must end with a newline - bulk_url = f"{base_url}/_bulk" - try: - response = requests.post(bulk_url, auth=auth, verify=verify, headers=headers, data=bulk_data) - if response.status_code in (200, 201): - result = response.json() - # Check for errors in result - if result.get('errors'): - for item in result['items']: - op = list(item.keys())[0] - if item[op].get('error'): - total_errors += 1 - self.console.log(f"[red]Error indexing document: {item[op]['error']}[/red]") - else: - total_indexed += 1 + docs = [json.loads(line.strip())] + if fn_convert: + docs = fn_convert(docs) + for doc in docs: + # Allow per-document _op_type (e.g., 'index', 'create', 'update', 'delete') + op_type = doc.pop('_op_type', 'index') + action_metadata = {"_index": index_name} + if '_id' in doc: + action_metadata['_id'] = doc.pop('_id') + action_line = {op_type: action_metadata} + action_line_json = json.dumps(action_line) + if op_type == 'delete': + # For delete operation, no source needed + actions.append(action_line_json) + action_count += 1 + elif op_type == 'update': + # For update operation, wrap the doc in 'doc' + doc_line_json = json.dumps({'doc': doc}) + actions.append(action_line_json) + actions.append(doc_line_json) + action_count += 1 + else: + # For index or create operation + doc_line_json = json.dumps(doc) + actions.append(action_line_json) + actions.append(doc_line_json) + action_count += 1 + + progress.update(task, advance=1) + + # Send batch to OpenSearch when batch_size is reached + if action_count >= batch_size: + # Prepare bulk data + bulk_data = '\n'.join(actions) + '\n' # Must end with a newline + bulk_url = f"{base_url}/_bulk" + try: + response = requests.post(bulk_url, auth=auth, verify=verify, headers=headers, data=bulk_data) + if response.status_code in (200, 201): + result = response.json() + # Check for errors in result + if result.get('errors'): + for item in result['items']: + op = list(item.keys())[0] + if item[op].get('error'): + total_errors += 1 + self.console.log(f"[red]Error indexing document: {item[op]['error']}[/red]") + else: + total_indexed += 1 + else: + total_indexed += len(result['items']) else: - total_indexed += len(result['items']) - else: - self.console.log(f"[red]Error indexing batch: {response.status_code} {response.text}[/red]") + self.console.log(f"[red]Error indexing batch: {response.status_code} {response.text}[/red]") + total_errors += action_count + except requests.exceptions.RequestException as e: + self.console.log(f"[red]Error communicating with OpenSearch during bulk indexing: {e}[/red]") total_errors += action_count - except requests.exceptions.RequestException as e: - self.console.log(f"[red]Error communicating with OpenSearch during bulk indexing: {e}[/red]") - total_errors += action_count - actions = [] # Clear the batch - action_count = 0 + actions = [] # Clear the batch + action_count = 0 except json.JSONDecodeError as e: self.console.log(f"[red]Failed to decode JSON:[/red] {e}") except Exception as e: @@ -243,7 +252,10 @@ class OpenSearchIndexer(BaseCommand): data_source.close() @OpenSearchIndexer.register -def index(self, index_name: str, batch_size: int = 1000, delete_first: bool = False, settings_file: str = None, mappings_file: str = None, file_path: str = None): +def index(self, index_name: str, batch_size: int = 1000, + delete_first: bool = False, + settings_file: str = None, mappings_file: str = None, + file_path: str = None, convert_fn: str = None, **kwargs): """ Pushes JSON data from stdin or a file to OpenSearch with batch indexing. @@ -289,6 +301,19 @@ def index(self, index_name: str, batch_size: int = 1000, delete_first: bool = Fa self.console.log(f"[red]Error loading mappings file '{mappings_file}':[/red] {e}") sys.exit(1) + if convert_fn is not None: + try: + # Split the class path into module and class name + module_path, class_name = convert_fn.rsplit('.', 1) + # Import the module + module = importlib.import_module(module_path) + cmd = getattr(module, class_name) + convert_fn = cmd(**kwargs) + except Exception as e: + raise ValueError("Invalid convert_fn argument. Must be a valid class specifiyer (e.g. owilix.plugins.push.convert.ConvertBase).", e) + self.console.print(f"Starting to push data to OpenSearch index '{index_name}'...") - self.push_to_opensearch(index_name, batch_size, delete_first, settings=settings, mappings=mappings, file_path=file_path) + self.push_to_opensearch(index_name, batch_size, delete_first, + settings=settings, mappings=mappings, + file_path=file_path, fn_convert=convert_fn) self.console.print("Finished pushing data to OpenSearch.") diff --git a/pyproject.toml b/pyproject.toml index 8c68390..2e73713 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -40,6 +40,7 @@ rich = "^13.7.0" numpy = ">2.0.0" typer = "^0.12.5" requests = "^2.32.3" +python-dateutil = "^2.9.0.post0" [[tool.poetry.source]] -- 2.51.2