From 594de1c79a4e6f10411dc9643f3c51a756aac0ad Mon Sep 17 00:00:00 2001 From: Nick Roberts Date: Mon, 23 Sep 2024 16:53:01 -0400 Subject: [PATCH] add support for version of doc upload which accepts bucket and JSON key --- stormcloud/ms/upload_documents.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/stormcloud/ms/upload_documents.py b/stormcloud/ms/upload_documents.py index 72378aa..b24f57e 100644 --- a/stormcloud/ms/upload_documents.py +++ b/stormcloud/ms/upload_documents.py @@ -92,7 +92,9 @@ def _sanitize(input_str: str) -> str: def get_ranked_docs(bucket: str, s3_client: Any, **kwargs) -> Iterator[dict]: - json_key = construct_key(**kwargs) + json_key = kwargs.get("json_key", None) + if json_key == None: + json_key = construct_key(**kwargs) logging.info(f"getting ranked document information from s3://{bucket}/{json_key}") res = s3_client.get_object(Bucket=bucket, Key=json_key) text = res.get("Body").read().decode() @@ -114,7 +116,7 @@ def __enter__(self): def __exit__(self, *args) -> None: self.handle_queue() - non_null_args = [a for a in filter(args)] + non_null_args = [a for a in filter(None, args)] if len(non_null_args) > 0: logging.error(f"exited with error message: {', '.join(non_null_args)}") @@ -153,9 +155,9 @@ def main(ms_client: Client, s3_client: Any, index: str, s3_bucket: str, update: import argparse import os - from dotenv import load_dotenv from client_utils import create_meilisearch_client, create_s3_client from constants import INDEX + from dotenv import load_dotenv load_dotenv("stormcloud/ms/.env")