Marmelator/pyinfra
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 36 Wiki Activity

adjust mock scriptis,add different file format option for IO

Browse Source
This commit is contained in:
Julius Unverfehrt 2022-03-17 12:46:51 +01:00
parent 3b50f88482
commit 6a36f56a3e
2 changed files with 35 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
scripts/manage_minio.py
View File

@ -37,8 +37,11 @@ def upload_compressed_response(storage, bucket_name, dossier_id, file_id, result
def add_file_compressed(storage, bucket_name, dossier_id, path) -> None:
path_gz = combine_dossier_id_and_file_id_and_extension(dossier_id, Path(path).stem, ".ORIGIN.pdf.gz")
if Path(path).suffix == ".pdf":
suffix_gz = ".ORIGIN.pdf.gz"
if Path(path).suffix == ".json":
suffix_gz = ".TEXT.json.gz"
path_gz = combine_dossier_id_and_file_id_and_extension(dossier_id, Path(path).stem, suffix_gz)
with open(path, "rb") as f:
data = gzip.compress(f.read())

43
scripts/mock_client.py
View File

@ -1,3 +1,4 @@
import argparse
import json
import pika
@ -6,6 +7,13 @@ from pyinfra.config import CONFIG
from pyinfra.storage.storages import get_s3_storage
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument("--analysis_container", "-a", choices=["detr", "ner", "image"], required=True)
args = parser.parse_args()
return args
def read_connection_params():
credentials = pika.PlainCredentials(CONFIG.rabbitmq.user, CONFIG.rabbitmq.password)
parameters = pika.ConnectionParameters(
@ -25,7 +33,7 @@ def make_channel(connection) -> pika.adapters.blocking_connection.BlockingChanne
def declare_queue(channel, queue: str):
args = {"x-dead-letter-exchange": "", "x-dead-letter-routing-key": CONFIG.rabbitmq.queues.dead_letter}
return channel.queue_declare(queue=queue, auto_delete=False, arguments=args)
return channel.queue_declare(queue=queue, auto_delete=False, durable=True, arguments=args)
def make_connection() -> pika.BlockingConnection:
@ -34,32 +42,41 @@ def make_connection() -> pika.BlockingConnection:
return connection
def build_message_bodies():
def update_message(message_dict, analyse_container_type):
if analyse_container_type == "detr" or analyse_container_type == "image":
message_dict.update({"targetFileExtension": "ORIGIN.pdf.gz", "responseFileExtension": "IMAGE_INFO.json.gz"})
if analyse_container_type == "ner":
message_dict.update({"targetFileExtension": "TEXT.json.gz", "responseFileExtension": "NER_ENTITIES.json.gz"})
return message_dict
def build_message_bodies(analyse_container_type):
storage = get_s3_storage()
for bucket_name, pdf_name in storage.get_all_object_names(CONFIG.storage.bucket):
file_id = pdf_name.split(".")[0]
dossier_id, file_id = file_id.split("/")
yield json.dumps(
{
"dossierId": dossier_id,
"fileId": file_id,
"targetFileExtension": "ORIGIN.pdf.gz",
"responseFileExtension": "detr.json.gz",
}
).encode()
message_dict = {
"dossierId": dossier_id,
"fileId": file_id
}
update_message(message_dict, analyse_container_type)
yield json.dumps(message_dict).encode()
if __name__ == "__main__":
def main(args):
connection = make_connection()
channel = make_channel(connection)
declare_queue(channel, CONFIG.rabbitmq.queues.input)
declare_queue(channel, CONFIG.rabbitmq.queues.output)
for body in build_message_bodies():
for body in build_message_bodies(args.analysis_container):
channel.basic_publish("", CONFIG.rabbitmq.queues.input, body)
print(f"Put {body} on {CONFIG.rabbitmq.queues.input}")
for method_frame, _, body in channel.consume(queue=CONFIG.rabbitmq.queues.output):
print(f"Received {json.loads(body)}")
channel.basic_ack(method_frame.delivery_tag)
if __name__ == "__main__":
main(parse_args())