Add Error handling Pipeline Execution (#1466)

Added Better Error Handling and additional test case to test out
different behaviours of the `pipeline.py`

Author: sundarshankar89

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

@@ -46,10 +46,14 @@ def _execute_sql_step(self, step: Step):
 
         # Execute the query using the database manager
         logging.info(f"Executing query: {query}")
-        result = self.executor.execute_query(query)
+        try:
+            result = self.executor.execute_query(query)
 
-        # Save the result to duckdb
-        self._save_to_db(result, step.name, str(step.mode))
+            # Save the result to duckdb
+            self._save_to_db(result, step.name, str(step.mode))
+        except Exception as e:
+            logging.error(f"SQL execution failed: {str(e)}")
+            raise RuntimeError(f"SQL execution failed: {str(e)}") from e
 
     def _execute_python_step(self, step: Step):
         logging.debug(f"Executing Python script: {step.extract_source}")

tests/integration/assessments/test_pipeline.py

@@ -22,12 +22,44 @@ def pipeline_config():
     return config
 
 
+@pytest.fixture(scope="module")
+def sql_failure_config():
+    prefix = Path(__file__).parent
+    config_path = f"{prefix}/../../resources/assessments/pipeline_config_sql_failure.yml"
+    config = PipelineClass.load_config_from_yaml(config_path)
+    for step in config.steps:
+        step.extract_source = f"{prefix}/../../{step.extract_source}"
+    return config
+
+
+@pytest.fixture(scope="module")
+def python_failure_config():
+    prefix = Path(__file__).parent
+    config_path = f"{prefix}/../../resources/assessments/pipeline_config_python_failure.yml"
+    config = PipelineClass.load_config_from_yaml(config_path)
+    for step in config.steps:
+        step.extract_source = f"{prefix}/../../{step.extract_source}"
+    return config
+
+
 def test_run_pipeline(extractor, pipeline_config, get_logger):
     pipeline = PipelineClass(config=pipeline_config, executor=extractor)
     pipeline.execute()
     assert verify_output(get_logger, pipeline_config.extract_folder)
 
 
+def test_run_sql_failure_pipeline(extractor, sql_failure_config, get_logger):
+    pipeline = PipelineClass(config=sql_failure_config, executor=extractor)
+    with pytest.raises(RuntimeError, match="SQL execution failed"):
+        pipeline.execute()
+
+
+def test_run_python_failure_pipeline(extractor, python_failure_config, get_logger):
+    pipeline = PipelineClass(config=python_failure_config, executor=extractor)
+    with pytest.raises(RuntimeError, match="Script execution failed"):
+        pipeline.execute()
+
+
 def verify_output(get_logger, path):
     conn = duckdb.connect(str(Path(path)) + "/" + DB_NAME)
 

tests/resources/assessments/invalid_query.sql

@@ -0,0 +1 @@
+SELECT * FROM non_existent_table;

tests/resources/assessments/invalid_script.py

@@ -0,0 +1,11 @@
+import sys
+import json
+
+
+def main():
+    print(json.dumps({"status": "error", "message": "This script is designed to fail"}), file=sys.stderr)
+    sys.exit(1)
+
+
+if __name__ == '__main__':
+    main()

tests/resources/assessments/pipeline_config_python_failure.yml

@@ -0,0 +1,10 @@
+name: "Python Failure Pipeline"
+version: "1.0"
+extract_folder: "tests/resources/assessments"
+steps:
+  - name: invalid_python_step
+    type: python
+    flag: active
+    extract_source: resources/assessments/invalid_script.py
+    mode: overwrite
+    frequency: daily

tests/resources/assessments/pipeline_config_sql_failure.yml

@@ -0,0 +1,10 @@
+name: "SQL Failure Pipeline"
+version: "1.0"
+extract_folder: /tmp/extracts/
+steps:
+  - name: invalid_sql_step
+    type: sql
+    flag: active
+    extract_source: resources/assessments/invalid_query.sql
+    mode: overwrite
+    frequency: daily