Merge pull request #7 from Nieuwe-Warmte-Nu/6-allow-resource-init-and-cleanup

6 allow resource init and cleanup

Author: MarkTNO

README.md

@@ -49,38 +49,42 @@ rabbitmq_config = RabbitmqConfig(
     "5678",
 )
 nwn_client = NwnClient(postgres_config, rabbitmq_config)
-job_id1: uuid4 = nwn_client.start_work_flow(
-    WorkFlowType.GROWTH_OPTIMIZER, "test_job1", "esdl_string", "test_user1", "test_proj"
-)
-job_id2: uuid4 = nwn_client.start_work_flow(
-    WorkFlowType.GROWTH_OPTIMIZER, "test_job2", "esdl_string", "test_user2", "test_proj"
-)
-print(job_id1)
-
-job1_input_esdl = nwn_client.get_job_input_esdl(job_id1)
-print(f"===== job1 input ESDL: {job1_input_esdl}")
-
-job1_status = nwn_client.get_job_status(job_id1)
-print(f"===== job1 status: {job1_status}")
-
-job1 = nwn_client.get_job_details(job_id1)
-print(f"===== {job1.job_name} input esdl: {job1.input_esdl}")
-
-jobs_all = nwn_client.get_all_jobs()
-print(f"===== {jobs_all[1].job_name} added at: {jobs_all[1].added_at}")
-
-
-jobs_from_ids = nwn_client.get_jobs_from_ids([job_id1, job_id2])
-print(f"===== {jobs_from_ids[1].job_name} added at: {jobs_from_ids[1].added_at}")
-
-
-jobs_from_user = nwn_client.get_jobs_from_user("test_user1")
-print(f"===== Jobs from test_user1 added at: {','.join([str(job.added_at) for job in jobs_from_user])}")
-
-
-jobs_from_project = nwn_client.get_jobs_from_project("test_proj")
-print(f"===== Jobs from test_proj added at: {','.join([str(job.added_at) for job in jobs_from_project])}")
-
-print(f"===== Deleted job with id '{job_id1}': {nwn_client.delete_job(job_id1)}")
+try:
+    nwn_client.connect()
+    job_id1: uuid4 = nwn_client.start_work_flow(
+        WorkFlowType.GROWTH_OPTIMIZER, "test_job1", "esdl_string", "test_user1", "test_proj"
+    )
+    job_id2: uuid4 = nwn_client.start_work_flow(
+        WorkFlowType.GROWTH_OPTIMIZER, "test_job2", "esdl_string", "test_user2", "test_proj"
+    )
+    print(job_id1)
+    
+    job1_input_esdl = nwn_client.get_job_input_esdl(job_id1)
+    print(f"===== job1 input ESDL: {job1_input_esdl}")
+    
+    job1_status = nwn_client.get_job_status(job_id1)
+    print(f"===== job1 status: {job1_status}")
+    
+    job1 = nwn_client.get_job_details(job_id1)
+    print(f"===== {job1.job_name} input esdl: {job1.input_esdl}")
+    
+    jobs_all = nwn_client.get_all_jobs()
+    print(f"===== {jobs_all[1].job_name} added at: {jobs_all[1].added_at}")
+    
+    
+    jobs_from_ids = nwn_client.get_jobs_from_ids([job_id1, job_id2])
+    print(f"===== {jobs_from_ids[1].job_name} added at: {jobs_from_ids[1].added_at}")
+    
+    
+    jobs_from_user = nwn_client.get_jobs_from_user("test_user1")
+    print(f"===== Jobs from test_user1 added at: {','.join([str(job.added_at) for job in jobs_from_user])}")
+    
+    
+    jobs_from_project = nwn_client.get_jobs_from_project("test_proj")
+    print(f"===== Jobs from test_proj added at: {','.join([str(job.added_at) for job in jobs_from_project])}")
+    
+    print(f"===== Deleted job with id '{job_id1}': {nwn_client.delete_job(job_id1)}")
+finally:
+    nwn_client.stop()
 
 ```

src/nwnsdk/nwn_client.py

@@ -1,9 +1,9 @@
 from dataclasses import dataclass
-from typing import List
+from typing import List, Callable, Dict
 from uuid import uuid4
 
 from nwnsdk.postgres.dbmodels import Job
-from nwnsdk.rabbitmq.rabbitmq_client import RabbitmqClient
+from nwnsdk.rabbitmq.rabbitmq_client import RabbitmqClient, Queue, PikaCallback
 
 import logging
 from nwnsdk.postgres.postgres_client import PostgresClient
@@ -12,65 +12,45 @@
 LOGGER = logging.getLogger("nwnsdk")
 
 
-class NwnClient:
+class NwnClient(PostgresClient, RabbitmqClient):
     rabbitmq_client: RabbitmqClient
     postgres_client: PostgresClient
     logger: logging.Logger
 
     def __init__(self, postgres_config: PostgresConfig, rabbitmq_config: RabbitmqConfig):
-        self.rabbitmq_client = RabbitmqClient(rabbitmq_config)
-        self.postgres_client = PostgresClient(postgres_config)
+        PostgresClient.__init__(self, postgres_config)
+        RabbitmqClient.__init__(self, rabbitmq_config)
+
+    def connect(self):
+        PostgresClient._connect_postgres(self)
+        RabbitmqClient._connect_rabbitmq(self)
+
+    def stop(self):
+        PostgresClient._close_postgres(self)
+        RabbitmqClient._stop_rabbitmq(self)
 
     def start_work_flow(
-            self, work_flow_type: WorkFlowType, job_name: str, esdl_str: str, user_name: str, project_name: str
+        self, work_flow_type: WorkFlowType, job_name: str, esdl_str: str, user_name: str, project_name: str
     ) -> uuid4:
         job_id: uuid4 = uuid4()
-        self.postgres_client.send_input(
+        PostgresClient._send_input(
+            self,
             job_id=job_id,
             job_name=job_name,
             work_flow_type=work_flow_type,
             esdl_str=esdl_str,
             user_name=user_name,
             project_name=project_name,
         )
-        self.rabbitmq_client.send_start_work_flow(job_id, work_flow_type)
+        RabbitmqClient._send_start_work_flow(self, job_id, work_flow_type)
 
         return job_id
 
-    def get_job_status(self, job_id: uuid4) -> JobStatus:
-        return self.postgres_client.get_job_status(job_id)
-
-    def get_job_input_esdl(self, job_id: uuid4) -> str:
-        return self.postgres_client.get_job_input_esdl(job_id)
-
-    def get_job_output_esdl(self, job_id: uuid4) -> str:
-        return self.postgres_client.get_job_output_esdl(job_id)
-
-    def get_job_logs(self, job_id: uuid4) -> str:
-        return self.postgres_client.get_job_logs(job_id)
-
     def get_job_details(self, job_id: uuid4) -> Job:
-        return self.postgres_client.get_job(job_id)
+        return self.get_job(job_id)
 
     def get_all_jobs(self) -> List[Job]:
-        return self.postgres_client.get_jobs()
+        return self.get_jobs()
 
     def get_jobs_from_ids(self, job_ids: List[uuid4]) -> List[Job]:
-        return self.postgres_client.get_jobs(job_ids)
-
-    def get_jobs_from_user(self, user_name: str) -> List[Job]:
-        return self.postgres_client.get_jobs_from_user(user_name)
-
-    def get_jobs_from_project(self, project_name: str) -> List[Job]:
-        return self.postgres_client.get_jobs_from_project(project_name)
-
-    def delete_job(self, job_id: uuid4) -> bool:
-        return self.postgres_client.delete_job(job_id)
-
-    @property
-    def db_client(self) -> PostgresClient:
-        return self.postgres_client
-
-    @property
-    def broker_client(self) -> RabbitmqClient:
-        return self.rabbitmq_client
+        return self.get_jobs(job_ids)

src/nwnsdk/postgres/postgres_client.py

@@ -6,6 +6,7 @@
 from sqlalchemy import select, update, delete
 from sqlalchemy.orm import Session
 from sqlalchemy.orm.strategy_options import load_only
+from sqlalchemy.engine import Engine
 
 from nwnsdk import PostgresConfig, WorkFlowType
 from nwnsdk.postgres.database import initialize_db, session_scope
@@ -29,17 +30,27 @@
 
 
 class PostgresClient:
+    db_config: PostgresConfig
+    engine: Engine
+
     def __init__(self, postgres_config: PostgresConfig):
-        initialize_db("nwn", postgres_config)
-
-    def send_input(
-            self,
-            job_id: uuid4,
-            job_name: str,
-            work_flow_type: WorkFlowType,
-            esdl_str: str,
-            user_name: str,
-            project_name: str,
+        self.db_config = postgres_config
+
+    def _connect_postgres(self):
+        self.engine = initialize_db("nwn", self.db_config)
+
+    def _close_postgres(self):
+        if self.engine:
+            self.engine.dispose()
+
+    def _send_input(
+        self,
+        job_id: uuid4,
+        job_name: str,
+        work_flow_type: WorkFlowType,
+        esdl_str: str,
+        user_name: str,
+        project_name: str,
     ) -> None:
         with session_scope() as session:
             new_job = Job(

src/nwnsdk/rabbitmq/rabbitmq_client.py

@@ -2,11 +2,11 @@
 
 import logging
 from enum import Enum
-import signal
 from typing import Callable, Dict
 from uuid import uuid4
 
 import pika
+import pika.exceptions
 import json
 
 from pika.adapters.blocking_connection import BlockingChannel
@@ -33,49 +33,74 @@ def from_workflow_type(workflow_type: WorkFlowType) -> "Queue":
 
 
 class RabbitmqClient:
+    rabbitmq_is_running: bool
+    rabbitmq_config: RabbitmqConfig
     rabbitmq_exchange: str
+    connection: pika.BlockingConnection
     channel: BlockingChannel
     queue: str
 
     def __init__(self, config: RabbitmqConfig):
+        self.rabbitmq_is_running = False
+        self.rabbitmq_config = config
         self.rabbitmq_exchange = config.exchange_name
 
+    def _connect_rabbitmq(self):
         # initialize rabbitmq connection
-        LOGGER.info("Connecting to RabbitMQ at %s:%s as user %s", config.host, config.port, config.user_name)
-        credentials = pika.PlainCredentials(config.user_name, config.password)
+        LOGGER.info(
+            "Connecting to RabbitMQ at %s:%s as user %s",
+            self.rabbitmq_config.host,
+            self.rabbitmq_config.port,
+            self.rabbitmq_config.user_name,
+        )
+        credentials = pika.PlainCredentials(self.rabbitmq_config.user_name, self.rabbitmq_config.password)
         parameters = pika.ConnectionParameters(
-            config.host, config.port, "/", credentials, heartbeat=3600, blocked_connection_timeout=3600
+            self.rabbitmq_config.host,
+            self.rabbitmq_config.port,
+            "/",
+            credentials,
+            heartbeat=3600,
+            blocked_connection_timeout=3600,
+            connection_attempts=10,
         )
 
-        connection = pika.BlockingConnection(parameters)
+        self.connection = pika.BlockingConnection(parameters)
 
-        self.channel = connection.channel()
+        self.channel = self.connection.channel()
         self.channel.basic_qos(prefetch_size=0, prefetch_count=1)
         self.channel.exchange_declare(exchange=self.rabbitmq_exchange, exchange_type="topic")
         self.queue = self.channel.queue_declare(Queue.StartWorkflowOptimizer.value, exclusive=False).method.queue
         self.channel.queue_bind(self.queue, self.rabbitmq_exchange, routing_key=Queue.StartWorkflowOptimizer.value)
         LOGGER.info("Connected to RabbitMQ")
 
-    def wait_for_data(self, callbacks: Dict[Queue, PikaCallback]):
-        for queue, callback in callbacks.items():
-            self.channel.basic_consume(queue=queue.value, on_message_callback=callback, auto_ack=False)
-
-        def stop(signal, frame):
-            LOGGER.info("Received signal %s. Stopping..", signal)
-            self.channel.stop_consuming()
-
-        signal.signal(signal.SIGINT, stop)
-        signal.signal(signal.SIGTERM, stop)
-
-        LOGGER.info("Waiting for input...")
-        self.channel.start_consuming()
-
-    def send_start_work_flow(self, job_id: uuid4, work_flow_type: WorkFlowType):
+    def wait_for_work(self, callbacks: Dict[Queue, PikaCallback]):
+        self.rabbitmq_is_running = True
+
+        while self.rabbitmq_is_running:
+            try:
+                for queue, callback in callbacks.items():
+                    self.channel.basic_consume(queue=queue.value, on_message_callback=callback, auto_ack=False)
+                LOGGER.info("Waiting for input...")
+                self.channel.start_consuming()
+            except pika.exceptions.ConnectionClosedByBroker as exc:
+                LOGGER.info('Connection was closed by broker. Reason: "%s". Shutting down...', exc.reply_text)
+            except pika.exceptions.AMQPConnectionError:
+                LOGGER.info("Connection was lost, retrying...")
+                self._connect_rabbitmq()
+
+    def _send_start_work_flow(self, job_id: uuid4, work_flow_type: WorkFlowType):
         # TODO convert to protobuf
         # TODO job_id converted to string for json
         body = json.dumps({"job_id": str(job_id)})
-        self.send_output(Queue.from_workflow_type(work_flow_type), body)
+        self._send_output(Queue.from_workflow_type(work_flow_type), body)
 
-    def send_output(self, queue: Queue, message: str):
+    def _send_output(self, queue: Queue, message: str):
         body: bytes = message.encode("utf-8")
         self.channel.basic_publish(exchange=self.rabbitmq_exchange, routing_key=queue.value, body=body)
+
+    def _stop_rabbitmq(self):
+        self.rabbitmq_is_running = False
+        if self.channel:
+            self.channel.stop_consuming()
+        if self.connection:
+            self.connection.close()