From 7cf7db9560ad6392cad3746da05212905b45728d Mon Sep 17 00:00:00 2001
From: Will Girten <will.girten@databricks.com>
Date: Fri, 2 May 2025 12:33:35 -0400
Subject: [PATCH] Add pipeline and step comments.

---
 .../remorph/assessments/profiler_config.py    |  2 +
 .../integration/assessments/test_pipeline.py  | 38 +++++++++++++++++++
 2 files changed, 40 insertions(+)

diff --git a/src/databricks/labs/remorph/assessments/profiler_config.py b/src/databricks/labs/remorph/assessments/profiler_config.py
index fe92ce111..31ce13772 100644
--- a/src/databricks/labs/remorph/assessments/profiler_config.py
+++ b/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)
diff --git a/tests/integration/assessments/test_pipeline.py b/tests/integration/assessments/test_pipeline.py
index eaf228c21..95252ea47 100644
--- a/tests/integration/assessments/test_pipeline.py
+++ b/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