From cf01bfde281dce6a7f4f01ed18b34cae5b0b2f74 Mon Sep 17 00:00:00 2001
From: Redza Rosli <91888036+redzarosliCognite@users.noreply.github.com>
Date: Wed, 20 Dec 2023 00:29:15 +0800
Subject: [PATCH] feat: custom success message

---
 cognite/extractorutils/base.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/cognite/extractorutils/base.py b/cognite/extractorutils/base.py
index edd64f14..dc8edbdd 100644
--- a/cognite/extractorutils/base.py
+++ b/cognite/extractorutils/base.py
@@ -88,6 +88,7 @@ def __init__(
         handle_interrupts: bool = True,
         reload_config_interval: Optional[int] = 300,
         reload_config_action: ReloadConfigAction = ReloadConfigAction.DO_NOTHING,
+        success_message: str = "Successful shutdown"
     ):
         self.name = name
         self.description = description
@@ -102,6 +103,7 @@ def __init__(
         self.handle_interrupts = handle_interrupts
         self.reload_config_interval = reload_config_interval
         self.reload_config_action = reload_config_action
+        self.success_message = success_message
 
         self.started = False
         self.configured_logger = False
@@ -208,7 +210,7 @@ def _report_success(self) -> None:
                 ExtractionPipelineRun(
                     extpipe_external_id=self.extraction_pipeline.external_id,
                     status="success",
-                    message="Successful shutdown",
+                    message=self.success_message,
                 )
             )