annotations/Azaion.Loader/main.pyx

import threading
from threading import Thread
import traceback
from credentials cimport Credentials
from remote_command cimport RemoteCommand, CommandType
from remote_command_handler cimport RemoteCommandHandler
from file_data cimport FileData, UploadFileData
from api_client cimport ApiClient

cdef class CommandProcessor:
    cdef RemoteCommandHandler remote_handler
    cdef ApiClient api_client
    cdef bint running
    cdef object shutdown_event
    cdef RemoteCommand ok_response

    def __init__(self, int zmq_port, str api_url):
        self.api_client = ApiClient(api_url)
        self.shutdown_event = threading.Event()
        self.remote_handler = RemoteCommandHandler(zmq_port, self.on_command)
        self.remote_handler.start()
        self.running = True
        self.ok_response = RemoteCommand(CommandType.OK)

    def start(self):
        while self.running:
            try:
                while not self.shutdown_event.is_set():
                    self.shutdown_event.wait(timeout=1.0)
            except Exception as e:
                traceback.print_exc()
        print('EXIT!')


    cdef on_command(self, RemoteCommand command):
        try:
            if command.command_type == CommandType.EXIT:
                self.remote_handler.send(command.client_id, self.ok_response.serialize())
                t = Thread(target=self.stop) # non-block worker:
                t.start()
                return

            if command.command_type == CommandType.LOGIN:
                self.api_client.set_credentials(Credentials.from_msgpack(command.data))
                self.remote_handler.send(command.client_id, self.ok_response.serialize())
            elif command.command_type == CommandType.LOAD:
                file_data = FileData.from_msgpack(command.data)
                file_bytes = self.api_client.load_bytes(file_data.filename, file_data.folder)
                self.remote_handler.send(command.client_id, RemoteCommand(CommandType.DATA_BYTES, file_bytes).serialize())
            elif command.command_type == CommandType.LOAD_BIG_SMALL:
                data = FileData.from_msgpack(command.data)
                file_bytes = self.api_client.load_big_small_resource(data.filename, data.folder)
                self.remote_handler.send(command.client_id, RemoteCommand(CommandType.DATA_BYTES, file_bytes).serialize())
            elif command.command_type == CommandType.UPLOAD_BIG_SMALL:
                data = UploadFileData.from_msgpack(command.data)
                file_bytes = self.api_client.upload_big_small_resource(data.resource, data.filename, data.folder)
                self.remote_handler.send(command.client_id, RemoteCommand(CommandType.OK).serialize())
            else:
                pass
        except Exception as e:
            print(f"Error handling client: {e}")
            err_command = RemoteCommand(CommandType.ERROR, None, str(e))
            self.remote_handler.send(command.client_id, err_command.serialize())

    def stop(self):
        self.shutdown_event.set()
        self.remote_handler.stop()
        self.running = False