Add Notification support

jupyter_scheduler/models.py

@@ -1,8 +1,8 @@
 import os
 from enum import Enum
-from typing import Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional, Type, Union
 
-from pydantic import BaseModel, root_validator
+from pydantic import BaseModel, root_validator, validator
 
 Tags = List[str]
 EnvironmentParameterValues = Union[int, float, bool, str]
@@ -11,6 +11,65 @@
 SCHEDULE_RE = ""
 
 
+class NotificationEvent(str, Enum):
+    """
+    Enum that represents events triggering notifications. Implementers can extend
+    this enum to include additional notification events as needed.
+
+    Attributes:
+        SUCCESS (str): Sent when a job completes successfully.
+        FAILURE (str): Sent on job failure.
+        STOPPED (str): Sent when a job is manually stopped.
+    """
+
+    SUCCESS = "Success"
+    FAILURE = "Failure"
+    STOPPED = "Stopped"
+
+
+class NotificationsConfig(BaseModel):
+    """Represents configuration for notifications.
+
+    Attributes:
+        send_to (List[str]): A list of symbols (e.g., email addresses) to which notifications should be sent.
+        events (List[NotificationEvent]): A list of events that should trigger the sending of notifications.
+        include_output (bool): A flag indicating whether a output should be included in the notification. Default is False.
+    """
+
+    send_to: List[str] = []
+    events: List[NotificationEvent] = []
+    include_output: bool = False
+
+    class Config:
+        orm_mode = True
+
+    @validator("send_to")
+    def validate_send_to(cls, v):
+        if len(v) > 100:
+            raise ValueError("Too many 'Send to' addressee identifiers. Maximum allowed is 100.")
+        return v
+
+    @validator("send_to", each_item=True)
+    def validate_send_to_items(cls, v):
+        if len(v) > 100:
+            raise ValueError(
+                "Each 'Send to' addressee identifier should be at most 100 characters long."
+            )
+        return v
+
+    @validator("events")
+    def validate_events(cls, v):
+        if len(v) > 100:
+            raise ValueError("Too many notification events. Maximum allowed is 100.")
+        return v
+
+    @validator("events", each_item=True)
+    def validate_events_items(cls, v):
+        if len(v.value) > 100:
+            raise ValueError("Each notification event should be at most 100 characters long.")
+        return v
+
+
 class RuntimeEnvironment(BaseModel):
     """Defines a runtime context where job
     execution will happen. For example, conda
@@ -26,6 +85,8 @@ class RuntimeEnvironment(BaseModel):
     compute_types: Optional[List[str]]
     default_compute_type: Optional[str]  # Should be a member of the compute_types list
     utc_only: Optional[bool]
+    notifications_enabled: bool = False
+    notification_events: List[Type[NotificationEvent]] = []
 
     def __str__(self):
         return self.json()
@@ -85,6 +146,7 @@ class CreateJob(BaseModel):
     name: str
     output_filename_template: Optional[str] = OUTPUT_FILENAME_TEMPLATE
     compute_type: Optional[str] = None
+    notifications_config: Optional[NotificationsConfig] = None
 
     @root_validator
     def compute_input_filename(cls, values) -> Dict:
@@ -145,6 +207,7 @@ class DescribeJob(BaseModel):
     status: Status = Status.CREATED
     status_message: Optional[str] = None
     downloaded: bool = False
+    notifications_config: Optional[NotificationsConfig] = None
 
     class Config:
         orm_mode = True
@@ -209,6 +272,7 @@ class CreateJobDefinition(BaseModel):
     compute_type: Optional[str] = None
     schedule: Optional[str] = None
     timezone: Optional[str] = None
+    notifications_config: Optional[NotificationsConfig] = None
 
     @root_validator
     def compute_input_filename(cls, values) -> Dict:
@@ -234,6 +298,7 @@ class DescribeJobDefinition(BaseModel):
     create_time: int
     update_time: int
     active: bool
+    notifications_config: Optional[NotificationsConfig] = None
 
     class Config:
         orm_mode = True
@@ -253,6 +318,7 @@ class UpdateJobDefinition(BaseModel):
     active: Optional[bool] = None
     compute_type: Optional[str] = None
     input_uri: Optional[str] = None
+    notifications_config: Optional[NotificationsConfig] = None
 
 
 class ListJobDefinitionsQuery(BaseModel):

jupyter_scheduler/orm.py

@@ -4,8 +4,14 @@
 from uuid import uuid4
 
 import sqlalchemy.types as types
-from sqlalchemy import Boolean, Column, Integer, String, create_engine
-from sqlalchemy.orm import declarative_base, declarative_mixin, registry, sessionmaker
+from sqlalchemy import Boolean, Column, ForeignKey, Integer, String, create_engine
+from sqlalchemy.orm import (
+    declarative_base,
+    declarative_mixin,
+    registry,
+    relationship,
+    sessionmaker,
+)
 
 from jupyter_scheduler.models import EmailNotifications, Status
 from jupyter_scheduler.utils import get_utc_timestamp
@@ -67,6 +73,14 @@ def process_result_value(self, value, dialect):
 mapper_registry = registry()
 
 
+class NotificationsConfigTable(Base):
+    __tablename__ = "notifications_config"
+    id = Column(String(36), primary_key=True, default=generate_uuid)
+    include_output = Column(Boolean, default=False)
+    send_to = Column(JsonType, nullable=False)
+    events = Column(JsonType, nullable=False)
+
+
 @declarative_mixin
 class CommonColumns:
     runtime_environment_name = Column(String(256), nullable=False)
@@ -98,6 +112,8 @@ class Job(CommonColumns, Base):
     url = Column(String(256), default=generate_jobs_url)
     pid = Column(Integer)
     idempotency_token = Column(String(256))
+    notifications_config_id = Column(String(36), ForeignKey("notifications_config.id"))
+    notifications_config = relationship("NotificationsConfigTable", lazy="joined")
 
 
 class JobDefinition(CommonColumns, Base):
@@ -108,6 +124,8 @@ class JobDefinition(CommonColumns, Base):
     url = Column(String(256), default=generate_job_definitions_url)
     create_time = Column(Integer, default=get_utc_timestamp)
     active = Column(Boolean, default=True)
+    notifications_config_id = Column(String(36), ForeignKey("notifications_config.id"))
+    notifications_config = relationship("NotificationsConfigTable", lazy="joined")
 
 
 def create_tables(db_url, drop_tables=False):

jupyter_scheduler/scheduler.py

@@ -38,7 +38,12 @@
     UpdateJob,
     UpdateJobDefinition,
 )
-from jupyter_scheduler.orm import Job, JobDefinition, create_session
+from jupyter_scheduler.orm import (
+    Job,
+    JobDefinition,
+    NotificationsConfigTable,
+    create_session,
+)
 from jupyter_scheduler.utils import create_output_directory, create_output_filename
 
 
@@ -396,7 +401,17 @@ def create_job(self, model: CreateJob) -> str:
             if not model.output_formats:
                 model.output_formats = []
 
-            job = Job(**model.dict(exclude_none=True, exclude={"input_uri"}))
+            orm_notifications_config = None
+            if model.notifications_config:
+                orm_notifications_config = NotificationsConfigTable(
+                    **model.notifications_config.dict()
+                )
+                session.add(orm_notifications_config)
+
+            job_data = model.dict(exclude={"input_uri", "notifications_config"})
+            job_data["notifications_config"] = orm_notifications_config
+
+            job = Job(**job_data)
             session.add(job)
             session.commit()
 
@@ -534,7 +549,20 @@ def create_job_definition(self, model: CreateJobDefinition) -> str:
             if not self.file_exists(model.input_uri):
                 raise InputUriError(model.input_uri)
 
-            job_definition = JobDefinition(**model.dict(exclude_none=True, exclude={"input_uri"}))
+            orm_notifications_config = None
+            if model.notifications_config:
+                orm_notifications_config = NotificationsConfigTable(
+                    **model.notifications_config.dict()
+                )
+                session.add(orm_notifications_config)
+                session.flush()
+
+            job_definition_data = model.dict(
+                exclude={"input_uri", "notifications_config"}, exclude_none=True
+            )
+            job_definition = JobDefinition(
+                **job_definition_data, notifications_config=orm_notifications_config
+            )
             session.add(job_definition)
             session.commit()
 

src/components/notifications-config-detail.tsx

@@ -0,0 +1,45 @@
+import React from 'react';
+
+import { Card, CardContent, Stack, FormLabel } from '@mui/material';
+import { useTranslator } from '../hooks';
+import { Scheduler } from '../handler';
+import { LabeledValue } from './labeled-value';
+
+type INotificationsConfigDetailProps = {
+  notificationsConfig: Scheduler.INotificationsConfig;
+};
+
+export function NotificationsConfigDetail(
+  props: INotificationsConfigDetailProps
+): JSX.Element {
+  const trans = useTranslator('jupyterlab');
+  const sendTo = props.notificationsConfig.send_to.join(', ');
+  const events = props.notificationsConfig.events.join(', ');
+
+  return (
+    <Card>
+      <CardContent>
+        <FormLabel component="legend" sx={{ mb: 2 }}>
+          {trans.__('Notifications Settings')}
+        </FormLabel>
+        <Stack spacing={2}>
+          <LabeledValue
+            label={trans.__('Send to')}
+            value={sendTo}
+            style={{ flex: '1 1 100%' }}
+          />
+          <LabeledValue
+            label={trans.__('Notification events')}
+            value={events}
+            style={{ flex: '1 1 100%' }}
+          />
+          <LabeledValue
+            label={trans.__('Include output')}
+            value={props.notificationsConfig.include_output ? 'True' : 'False'}
+            style={{ flex: '1 1 100%' }}
+          />
+        </Stack>
+      </CardContent>
+    </Card>
+  );
+}