refactor: address review comments

Author: stevenhsd

src/dve/common/error_utils.py

@@ -13,23 +13,23 @@
 from dve.core_engine.exceptions import CriticalProcessingError
 from dve.core_engine.loggers import get_logger
 from dve.core_engine.message import UserMessage
-from dve.core_engine.type_hints import URI, DVEStage, Messages
+from dve.core_engine.type_hints import URI, DVEStageName, Messages
 
 
-def get_feedback_errors_uri(working_folder: URI, step_name: DVEStage) -> URI:
+def get_feedback_errors_uri(working_folder: URI, step_name: DVEStageName) -> URI:
     """Determine the location of json lines file containing all errors generated in a step"""
     return fh.joinuri(working_folder, "errors", f"{step_name}_errors.jsonl")
 
 
 def get_processing_errors_uri(working_folder: URI) -> URI:
     """Determine the location of json lines file containing all processing
     errors generated from DVE run"""
-    return fh.joinuri(working_folder, "errors", "processing_errors", "processing_errors.jsonl")
+    return fh.joinuri(working_folder, "processing_errors", "processing_errors.jsonl")
 
 
 def dump_feedback_errors(
     working_folder: URI,
-    step_name: DVEStage,
+    step_name: DVEStageName,
     messages: Messages,
     key_fields: Optional[dict[str, list[str]]] = None,
 ) -> URI:
@@ -76,7 +76,7 @@ def dump_processing_errors(
     if not errors:
         raise AttributeError("errors list not passed")
 
-    error_file: URI = fh.joinuri(working_folder, "processing_errors", "processing_errors.json")
+    error_file: URI = get_processing_errors_uri(working_folder)
     processed = []
 
     for error in errors:
@@ -121,7 +121,7 @@ class BackgroundMessageWriter:
     def __init__(
         self,
         working_directory: URI,
-        dve_stage: DVEStage,
+        dve_stage: DVEStageName,
         key_fields: Optional[dict[str, list[str]]] = None,
         logger: Optional[logging.Logger] = None,
     ):
@@ -149,6 +149,7 @@ def write_thread(self) -> Thread:  # type: ignore
 
     def _write_process_wrapper(self):
         """Wrapper for dump feedback errors to run in background process"""
+        # writing thread will block if nothing in queue
         while True:
             if msgs := self.write_queue.get():
                 dump_feedback_errors(
@@ -167,6 +168,7 @@ def __exit__(self, exc_type, exc_value, traceback):
                 "Issue occured during background write process:",
                 exc_info=(exc_type, exc_value, traceback),
             )
+        # None value in queue will trigger break in target
         self.write_queue.put(None)
         self.write_thread.join()
 

src/dve/core_engine/backends/base/contract.py

@@ -27,6 +27,7 @@
 from dve.core_engine.type_hints import (
     URI,
     ArbitraryFunction,
+    DVEStageName,
     EntityLocations,
     EntityName,
     JSONDict,
@@ -102,6 +103,10 @@ class BaseDataContract(Generic[EntityType], ABC):
     This is set and populated in `__init_subclass__` by identifying methods
     decorated with the '@reader_override' decorator, and is used in `read_entity_type`.
 
+    """
+    __stage_name__: DVEStageName = "data_contract"
+    """
+    The name of the data contract DVE stage for use in auditing and logging
     """
 
     def __init_subclass__(cls, *_, **__) -> None:
@@ -392,13 +397,13 @@ def apply(
         and return the validated entities and any messages.
 
         """
-        feedback_errors_uri = get_feedback_errors_uri(working_dir, "data_contract")
+        feedback_errors_uri = get_feedback_errors_uri(working_dir, self.__stage_name__)
         processing_errors_uri = get_processing_errors_uri(working_dir)
         entities, messages, successful = self.read_raw_entities(entity_locations, contract_metadata)
         if not successful:
             dump_processing_errors(
                 working_dir,
-                "data_contract",
+                self.__stage_name__,
                 [
                     CriticalProcessingError(
                         "Issue occurred while reading raw entities",
@@ -416,15 +421,15 @@ def apply(
             successful = False
             new_messages = render_error(
                 err,
-                "data contract",
+                self.__stage_name__,
                 self.logger,
             )
             dump_processing_errors(
                 working_dir,
-                "data_contract",
+                self.__stage_name__,
                 [
                     CriticalProcessingError(
-                        "Issue occurred while applying data_contract",
+                        f"Issue occurred while applying {self.__stage_name__}",
                         [msg.error_message for msg in new_messages],
                     )
                 ],

src/dve/core_engine/backends/base/rules.py

@@ -45,7 +45,7 @@
 from dve.core_engine.backends.types import Entities, EntityType, StageSuccessful
 from dve.core_engine.exceptions import CriticalProcessingError
 from dve.core_engine.loggers import get_logger
-from dve.core_engine.type_hints import URI, EntityName, Messages, TemplateVariables
+from dve.core_engine.type_hints import URI, DVEStageName, EntityName, Messages, TemplateVariables
 
 T_contra = TypeVar("T_contra", bound=AbstractStep, contravariant=True)
 T = TypeVar("T", bound=AbstractStep)
@@ -88,6 +88,10 @@ class BaseStepImplementations(Generic[EntityType], ABC):  # pylint: disable=too-
 
     This will be populated from the generic annotation at class creation time.
 
+    """
+    __stage_name__: DVEStageName = "business_rules"
+    """
+    The name of the business rules DVE stage for use in auditing and logging
     """
 
     def __init_subclass__(cls, *_, **__) -> None:
@@ -367,13 +371,13 @@ def apply_sync_filters(
 
         """
         filters_by_entity: dict[EntityName, list[DeferredFilter]] = defaultdict(list)
-        feedback_errors_uri = get_feedback_errors_uri(working_directory, "business_rules")
+        feedback_errors_uri = get_feedback_errors_uri(working_directory, self.__stage_name__)
         for rule in filters:
             filters_by_entity[rule.entity_name].append(rule)
 
         with BackgroundMessageWriter(
             working_directory=working_directory,
-            dve_stage="business_rules",
+            dve_stage=self.__stage_name__,
             key_fields=key_fields,
             logger=self.logger,
         ) as msg_writer:
@@ -402,7 +406,7 @@ def apply_sync_filters(
                         if not success:
                             processing_errors_uri = dump_processing_errors(
                                 working_directory,
-                                "business_rules",
+                                self.__stage_name__,
                                 [
                                     CriticalProcessingError(
                                         "Issue occurred while applying filter logic",
@@ -431,7 +435,7 @@ def apply_sync_filters(
                         if not success:
                             processing_errors_uri = dump_processing_errors(
                                 working_directory,
-                                "business_rules",
+                                self.__stage_name__,
                                 [
                                     CriticalProcessingError(
                                         "Issue occurred while generating FeedbackMessages",
@@ -459,7 +463,7 @@ def apply_sync_filters(
                         if not success:
                             processing_errors_uri = dump_processing_errors(
                                 working_directory,
-                                "business_rules",
+                                self.__stage_name__,
                                 [
                                     CriticalProcessingError(
                                         "Issue occurred while generating FeedbackMessages",
@@ -497,7 +501,7 @@ def apply_sync_filters(
                     if not success:
                         processing_errors_uri = dump_processing_errors(
                             working_directory,
-                            "business_rules",
+                            self.__stage_name__,
                             [
                                 CriticalProcessingError(
                                     "Issue occurred while filtering error records",
@@ -525,7 +529,7 @@ def apply_sync_filters(
                         if not success:
                             processing_errors_uri = dump_processing_errors(
                                 working_directory,
-                                "business_rules",
+                                self.__stage_name__,
                                 [
                                     CriticalProcessingError(
                                         "Issue occurred while generating FeedbackMessages",
@@ -555,7 +559,7 @@ def apply_rules(
         """
         self.logger.info("Applying business rules")
         rules_and_locals: Iterable[tuple[Rule, TemplateVariables]]
-        errors_uri = get_feedback_errors_uri(working_directory, "business_rules")
+        errors_uri = get_feedback_errors_uri(working_directory, self.__stage_name__)
         if rule_metadata.templating_strategy == "upfront":
             rules_and_locals = []
             for rule, local_variables in rule_metadata:
@@ -584,7 +588,7 @@ def apply_rules(
                 if not success:
                     processing_errors_uri = dump_processing_errors(
                         working_directory,
-                        "business_rules",
+                        self.__stage_name__,
                         [
                             CriticalProcessingError(
                                 "Issue occurred while applying pre filter steps",
@@ -593,14 +597,16 @@ def apply_rules(
                         ],
                     )
                     if pre_sync_messages:
-                        dump_feedback_errors(working_directory, "business_rules", pre_sync_messages)
+                        dump_feedback_errors(
+                            working_directory, self.__stage_name__, pre_sync_messages
+                        )
 
                     return processing_errors_uri, False
                 # if not a failure, ensure we keep track of any informational messages
                 pre_sync_messages.extend(stage_messages)
             # if all successful, ensure we write out all informational messages
             if pre_sync_messages:
-                dump_feedback_errors(working_directory, "business_rules", pre_sync_messages)
+                dump_feedback_errors(working_directory, self.__stage_name__, pre_sync_messages)
 
         sync_steps = []
         for rule, local_variables in rules_and_locals:
@@ -634,7 +640,7 @@ def apply_rules(
                 if not success:
                     processing_errors_uri = dump_processing_errors(
                         working_directory,
-                        "business_rules",
+                        self.__stage_name__,
                         [
                             CriticalProcessingError(
                                 "Issue occurred while applying post filter steps",
@@ -644,15 +650,15 @@ def apply_rules(
                     )
                     if post_sync_messages:
                         dump_feedback_errors(
-                            working_directory, "business_rules", post_sync_messages
+                            working_directory, self.__stage_name__, post_sync_messages
                         )
 
                     return processing_errors_uri, False
                 # if not a failure, ensure we keep track of any informational messages
                 post_sync_messages.extend(stage_messages)
             # if all successful, ensure we write out all informational messages
             if post_sync_messages:
-                dump_feedback_errors(working_directory, "business_rules", post_sync_messages)
+                dump_feedback_errors(working_directory, self.__stage_name__, post_sync_messages)
         return errors_uri, True
 
     def read_parquet(self, path: URI, **kwargs) -> EntityType:

src/dve/core_engine/backends/implementations/duckdb/contract.py

@@ -4,7 +4,7 @@
 import logging
 from collections.abc import Iterator
 from functools import partial
-from multiprocessing import Pool
+from multiprocessing import Pool, cpu_count
 from typing import Any, Optional
 from uuid import uuid4
 
@@ -164,7 +164,7 @@ def apply_data_contract(
 
                 batches = pq.ParquetFile(entity_locations[entity_name]).iter_batches(10000)
                 msg_count = 0
-                with Pool(8) as pool:
+                with Pool(cpu_count() - 1) as pool:
                     for msgs in pool.imap_unordered(row_validator_helper, batches):
                         if msgs:
                             msg_writer.write_queue.put(msgs)

src/dve/core_engine/type_hints.py

@@ -245,7 +245,7 @@
 BinaryComparator = Callable[[Any, Any], bool]
 """Type hint for operator functions"""
 
-DVEStage = Literal[
+DVEStageName = Literal[
     "audit_received",
     "file_transformation",
     "data_contract",

src/dve/pipeline/pipeline.py

@@ -35,7 +35,7 @@
 from dve.core_engine.loggers import get_logger
 from dve.core_engine.message import FeedbackMessage
 from dve.core_engine.models import SubmissionInfo, SubmissionStatisticsRecord
-from dve.core_engine.type_hints import URI, DVEStage, FileURI, InfoURI
+from dve.core_engine.type_hints import URI, DVEStageName, FileURI, InfoURI
 from dve.parser import file_handling as fh
 from dve.parser.file_handling.implementations.file import LocalFilesystemImplementation
 from dve.parser.file_handling.service import _get_implementation
@@ -113,7 +113,9 @@ def get_entity_count(entity: EntityType) -> int:
         """Get a row count of an entity stored as parquet"""
         raise NotImplementedError()
 
-    def get_submission_status(self, step_name: DVEStage, submission_id: str) -> SubmissionStatus:
+    def get_submission_status(
+        self, step_name: DVEStageName, submission_id: str
+    ) -> SubmissionStatus:
         """Determine submission status of a submission if not explicitly given"""
         if not (submission_status := self._audit_tables.get_submission_status(submission_id)):
             self._logger.warning(

tests/test_pipeline/test_foundry_ddb_pipeline.py

@@ -121,7 +121,7 @@ def test_foundry_runner_error(planet_test_files, temp_ddb_conn):
             processing_folder,
             sub_info.submission_id,
             "processing_errors",
-            "processing_errors.json"
+            "processing_errors.jsonl"
         )
         assert perror_path.exists()
         perror_schema = {

tests/test_reporting/test_error_utils.py

@@ -44,7 +44,7 @@ def test_dump_processing_errors():
             perror_schema
         )
         error_df = pl.read_json(
-            Path(output_path, "processing_errors.json")
+            Path(output_path, "processing_errors.jsonl")
         )
         cols_to_check = ["step_name", "error_location", "error_level", "error_message"]