Merge pull request #5 from Nieuwe-Warmte-Nu/add_additional_func

Add additional func

Author: lfse-slafleur

README.md

@@ -4,7 +4,13 @@ Nieuwe Warmte Nu
 
 ## install/update pip package
 see
-https://packaging.python.org/en/latest/tutorials/packaging-projects/#:~:text=Generating%20distribution%20archives
+https://packaging.python.org/en/latest/tutorials/packaging-projects/#:~:text=Generating%20distribution%20archives  
+Update version in `pyproject.toml`.  
+In root directory:
+```
+py -m build
+twine upload dist/*
+```
 
 ## usage
 Install in development mode, in root directory:
@@ -17,12 +23,13 @@ or install from pypi:
 pip install nwnsdk
 ```
 
-Start local rabbitmq and postgres, in root directory:
+For testing locally with RabbitMQ and Postgres on docker desktop clone https://github.com/Nieuwe-Warmte-Nu/computation-engine, copy `.env-template` to `.env` and in the root directory:
 ```
 docker-compose up
 ```
 
 Example usage
+
 ```python
 from uuid import uuid4
 from nwnsdk import NwnClient, WorkFlowType, PostgresConfig, RabbitmqConfig
@@ -42,7 +49,38 @@ rabbitmq_config = RabbitmqConfig(
     "5678",
 )
 nwn_client = NwnClient(postgres_config, rabbitmq_config)
-job_id: uuid4 = nwn_client.start_work_flow(WorkFlowType.GROWTH_OPTIMIZER, "test_job", "esdl_string", "test_user")
+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)}")
 
-print(job_id)
 ```

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "nwnsdk"
-version = "0.0.3"
+version = "0.0.4"
 authors = [{ name = "NieuweWarmteNu", email = "[email protected]" }]
 description = "Nieuwe Warmte Nu - compute engine sdk python"
 classifiers = [
@@ -20,10 +20,11 @@ readme = "README.md"
 requires-python = ">=3.10"
 
 dependencies = [
-    "python-dotenv == 1.0.0",
-    "pika == 1.3.2",
+    "python-dotenv==1.0.0",
+    "pika==1.3.2",
     "pika-stubs==0.1.3",
     "SQLAlchemy==2.0.19",
+    "sqlalchemy-stubs==0.4",
     "psycopg2-binary==2.9.7",
 ]
 

src/nwnsdk/app_logging.py

@@ -70,6 +70,8 @@ def setup_logging(log_level: LogLevel, logger_name: Optional[str]) -> logging.Lo
     ----------
     log_level : LogLevel
         The LogLevel for this logger.
+    logger_name : Optional[str]
+        Name for this logger.
     """
     logger = logging.getLogger(logger_name)
 

src/nwnsdk/nwn_client.py

@@ -1,10 +1,13 @@
+from dataclasses import dataclass
+from typing import List
 from uuid import uuid4
 
+from nwnsdk.postgres.dbmodels import Job
 from nwnsdk.rabbitmq.rabbitmq_client import RabbitmqClient
 
 import logging
 from nwnsdk.postgres.postgres_client import PostgresClient
-from nwnsdk import PostgresConfig, RabbitmqConfig, WorkFlowType
+from nwnsdk import PostgresConfig, RabbitmqConfig, WorkFlowType, JobStatus
 
 LOGGER = logging.getLogger("nwnsdk")
 
@@ -18,15 +21,52 @@ def __init__(self, postgres_config: PostgresConfig, rabbitmq_config: RabbitmqCon
         self.rabbitmq_client = RabbitmqClient(rabbitmq_config)
         self.postgres_client = PostgresClient(postgres_config)
 
-    def start_work_flow(self, work_flow_type: WorkFlowType, job_name: str, esdl_str: str, user_name: str) -> uuid4:
+    def start_work_flow(
+            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(
-            job_id=job_id, job_name=job_name, work_flow_type=work_flow_type, esdl_str=esdl_str, user_name=user_name
+            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)
 
         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)
+
+    def get_all_jobs(self) -> List[Job]:
+        return self.postgres_client.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

src/nwnsdk/postgres/database.py

@@ -15,19 +15,20 @@
 
 
 @contextmanager
-def session_scope(bind=None) -> Generator[SQLSession, None, None]:
+def session_scope(do_expunge=False) -> Generator[SQLSession, None, None]:
     """Provide a transactional scope around a series of operations. Ensures that the session is
     committed and closed. Exceptions raised within the 'with' block using this contextmanager
     should be handled in the with block itself. They will not be caught by the 'except' here."""
     try:
-        if bind:
-            yield Session(bind=bind)
         yield Session()
+
+        if do_expunge:
+            Session.expunge_all()
         Session.commit()
-    except Exception:
+    except Exception as e:
         # Only the exceptions raised by session.commit above are caught here
         Session.rollback()
-        raise
+        raise e
     finally:
         Session.remove()
 

src/nwnsdk/postgres/dbmodels.py

@@ -1,8 +1,11 @@
+from dataclasses import dataclass
 from enum import Enum
 
 import sqlalchemy as db
-
 from sqlalchemy.ext.declarative import declarative_base
+from sqlalchemy.dialects.postgresql import UUID
+
+from nwnsdk import WorkFlowType
 
 Base = declarative_base()
 
@@ -15,18 +18,20 @@ class JobStatus(Enum):
     STOPPED = "stopped"
 
 
+@dataclass
 class Job(Base):
     __tablename__ = "job"
 
-    job_id = db.Column(db.UUID, primary_key=True)
+    job_id = db.Column(UUID(as_uuid=True), primary_key=True)
     job_name = db.Column(db.String, nullable=False)
-    work_flow_type = db.Column(db.String, nullable=False)
-    map_editor_user = db.Column(db.String)
-    status = db.Column(db.String, nullable=False)
+    work_flow_type = db.Column(db.Enum(WorkFlowType), nullable=False)
+    user_name = db.Column(db.String)
+    project_name = db.Column(db.String)
+    status = db.Column(db.Enum(JobStatus), nullable=False)
     input_config = db.Column(db.String)
     input_esdl = db.Column(db.String, nullable=False)
     output_esdl = db.Column(db.String)
     added_at = db.Column(db.DateTime(timezone=True), nullable=False)
     running_at = db.Column(db.DateTime(timezone=True))
     stopped_at = db.Column(db.DateTime(timezone=True))
-    logs = db.Column(db.DateTime(timezone=True))
+    logs = db.Column(db.String)