Add Pipeline and Step Comment Attributes

src/databricks/labs/remorph/assessments/profiler_config.py

@@ -9,6 +9,7 @@ class Step:
     mode: str | None
     frequency: str | None
     flag: str | None
+    comment: str | None = None
 
     def __post_init__(self):
         if self.frequency is None:
@@ -24,4 +25,5 @@ class PipelineConfig:
     name: str
     version: str
     extract_folder: str
+    comment: str | None = None
     steps: list[Step] = field(default_factory=list)

tests/integration/assessments/test_pipeline.py

@@ -3,6 +3,8 @@
 import pytest
 
 from databricks.labs.remorph.assessments.pipeline import PipelineClass, DB_NAME
+
+from databricks.labs.remorph.assessments.profiler_config import Step, PipelineConfig
 from ..connections.helpers import get_db_manager
 
 
@@ -79,3 +81,39 @@ def verify_output(get_logger, path):
     conn.close()
     logger.info("All expected tables exist and are not empty")
     return True
+
+
+def test_pipeline_config_comments():
+    pipeline_w_comments = PipelineConfig(
+        name="warehouse_profiler",
+        version="1.0",
+        extract_folder="/tmp/extracts",
+        comment="A pipeline for extracting warehouse usage.",
+    )
+    pipeline_wo_comments = PipelineConfig(
+        name="another_warehouse_profiler", version="1.0", extract_folder="/tmp/extracts"
+    )
+    assert pipeline_w_comments.comment == "A pipeline for extracting warehouse usage."
+    assert pipeline_wo_comments.comment is None
+
+
+def test_pipeline_step_comments():
+    step_w_comment = Step(
+        name="step_w_comment",
+        type="sql",
+        extract_source="path/to/extract/source.sql",
+        mode="append",
+        frequency="once",
+        flag="active",
+        comment="This is a step comment.",
+    )
+    step_wo_comment = Step(
+        name="step_wo_comment",
+        type="python",
+        extract_source="path/to/extract/source.py",
+        mode="overwrite",
+        frequency="daily",
+        flag="inactive",
+    )
+    assert step_w_comment.comment == "This is a step comment."
+    assert step_wo_comment.comment is None