From b9c872e6f1df286cc451031676f7f2458efb2d2e Mon Sep 17 00:00:00 2001 From: stevenhsd <56357022+stevenhsd@users.noreply.github.com> Date: Fri, 25 Sep 2026 23:04:07 +0100 Subject: [PATCH 1/2] feat: add facility to add additional checks at end of file transformation stage --- src/dve/core_engine/backends/base/reader.py | 56 ++++- src/dve/core_engine/backends/base/rules.py | 8 +- .../implementations/duckdb/duckdb_helpers.py | 10 + .../implementations/duckdb/readers/csv.py | 2 + .../implementations/duckdb/readers/json.py | 2 + .../implementations/duckdb/readers/xml.py | 6 +- .../implementations/spark/readers/csv.py | 2 + .../implementations/spark/readers/json.py | 2 + .../implementations/spark/readers/xml.py | 3 + .../implementations/spark/spark_helpers.py | 10 + .../core_engine/configuration/v1/__init__.py | 11 + src/dve/pipeline/pipeline.py | 5 +- tests/features/flights.feature | 14 ++ tests/features/steps/steps_pipeline.py | 2 +- tests/features/steps/utilities.py | 2 +- .../flights_add_reader_checks.dischema.json | 208 ++++++++++++++++++ 16 files changed, 322 insertions(+), 21 deletions(-) create mode 100644 tests/testdata/flights/flights_add_reader_checks.dischema.json diff --git a/src/dve/core_engine/backends/base/reader.py b/src/dve/core_engine/backends/base/reader.py index ae0e99f..e65fb76 100644 --- a/src/dve/core_engine/backends/base/reader.py +++ b/src/dve/core_engine/backends/base/reader.py @@ -10,6 +10,7 @@ from dve.core_engine.backends.exceptions import MessageBearingError, ReaderLacksEntityTypeSupport from dve.core_engine.backends.types import EntityName, EntityType +from dve.core_engine.configuration.v1 import AllowedAdditionalChecks, _ReaderAdditionalChecksConfig from dve.core_engine.message import FeedbackMessage from dve.core_engine.type_hints import URI, ArbitraryFunction, WrapDecorator from dve.parser.file_handling.service import open_stream @@ -109,7 +110,10 @@ def read_to_entity_type( entity_name: EntityName, schema: type[BaseModel], all_model_fields: Optional[set[str]] = None, - ) -> EntityType: + additional_checks: Optional[ + dict[AllowedAdditionalChecks, _ReaderAdditionalChecksConfig] + ] = None, + ): """Read to the specified entity type, if supported. NOTE: Simple types should either be returned as strings (if present) or @@ -117,21 +121,47 @@ def read_to_entity_type( data contract. """ - if entity_name == Iterator[dict[str, Any]]: - return self.read_to_py_iterator( + additional_checks = additional_checks or {} + + self.raise_if_not_sensible_file(resource, entity_name) + + if entity_type == Iterator[dict[str, Any]]: + entity = self.read_to_py_iterator( resource, entity_name, schema, all_model_fields # type: ignore ) - self.raise_if_not_sensible_file(resource, entity_name) + else: - try: - reader_func = self.__read_methods__[entity_type] - except KeyError as err: - raise ReaderLacksEntityTypeSupport(entity_type=entity_type) from err + try: + reader_func = self.__read_methods__[entity_type] + except KeyError as err: + raise ReaderLacksEntityTypeSupport(entity_type=entity_type) from err - return reader_func( - self, resource, entity_name, schema, all_model_fields=all_model_fields # type: ignore - ) + entity = reader_func( + self, + resource, + entity_name, + schema, + all_model_fields=all_model_fields, # type: ignore + ) + + if config := additional_checks.get("check_empty"): + if self.check_entity_empty(entity): + raise MessageBearingError( + f"The mandatory entity {entity_name} is empty", + messages=[ + FeedbackMessage( + entity=entity_name, + record=None, + failure_type="submission", + error_location=entity_name, + error_code=config.error_code, + error_message=config.error_message, + ) + ], + ) + + return entity def add_record_index(self, entity: EntityType, **kwargs) -> EntityType: """Add a record index to the entity""" @@ -141,6 +171,10 @@ def drop_record_index(self, entity: EntityType, **kwargs) -> EntityType: """Drop a record index to the entity""" raise NotImplementedError(f"drop_record_index not implemented in {self.__class__}") + def check_entity_empty(self, entity: EntityType) -> bool: + """Determine if the entity supplied is empty""" + raise NotImplementedError(f"check_entity_empty not implemented in {self.__class__}") + def write_parquet( self, entity: EntityType, diff --git a/src/dve/core_engine/backends/base/rules.py b/src/dve/core_engine/backends/base/rules.py index ee93106..0b5b8a0 100644 --- a/src/dve/core_engine/backends/base/rules.py +++ b/src/dve/core_engine/backends/base/rules.py @@ -448,8 +448,7 @@ def process_node(node: HierarchyNode): record=record, # type: ignore error_location=location, error_message=template_object( - node.missing_parent_id_error_message, - record + node.missing_parent_id_error_message, record ), failure_type="record", error_type="record", @@ -522,10 +521,7 @@ def process_node(node: HierarchyNode) -> bool: entity=node.parent_entity, record=record, # type: ignore error_location=location, - error_message=template_object( - node.no_valid_records_error_message, - record - ), + error_message=template_object(node.no_valid_records_error_message, record), failure_type="record", error_type="record", error_code=node.no_valid_records_error_code, diff --git a/src/dve/core_engine/backends/implementations/duckdb/duckdb_helpers.py b/src/dve/core_engine/backends/implementations/duckdb/duckdb_helpers.py index 588cd7e..990f6f4 100644 --- a/src/dve/core_engine/backends/implementations/duckdb/duckdb_helpers.py +++ b/src/dve/core_engine/backends/implementations/duckdb/duckdb_helpers.py @@ -322,6 +322,16 @@ def duckdb_get_entity_count(cls): return cls +def _duckdb_check_entity_empty(self, entity: DuckDBPyRelation) -> bool: # pylint: disable=W0613 + return entity.shape[0] == 0 + + +def duckdb_check_entity_empty(cls): + """Class decorator to check whether a supplied entity is empty""" + cls.check_entity_empty = _duckdb_check_entity_empty + return cls + + def get_all_registered_udfs(connection: DuckDBPyConnection) -> set[str]: """Function to supply the names of a registered functions stored in the supplied duckdb connection. Creates the temp table used to store registered functions (if not exists). diff --git a/src/dve/core_engine/backends/implementations/duckdb/readers/csv.py b/src/dve/core_engine/backends/implementations/duckdb/readers/csv.py index 723e5e3..012673a 100644 --- a/src/dve/core_engine/backends/implementations/duckdb/readers/csv.py +++ b/src/dve/core_engine/backends/implementations/duckdb/readers/csv.py @@ -22,6 +22,7 @@ UnableToParseCSVError, ) from dve.core_engine.backends.implementations.duckdb.duckdb_helpers import ( + duckdb_check_entity_empty, duckdb_record_index, duckdb_write_parquet, get_duckdb_type_from_annotation, @@ -35,6 +36,7 @@ from dve.parser.file_handling import get_content_length +@duckdb_check_entity_empty @duckdb_record_index @duckdb_write_parquet class DuckDBCSVReader(CSVFileReader): diff --git a/src/dve/core_engine/backends/implementations/duckdb/readers/json.py b/src/dve/core_engine/backends/implementations/duckdb/readers/json.py index 79d74c6..84b601d 100644 --- a/src/dve/core_engine/backends/implementations/duckdb/readers/json.py +++ b/src/dve/core_engine/backends/implementations/duckdb/readers/json.py @@ -10,6 +10,7 @@ from dve.core_engine.backends.base.reader import BaseFileReader, read_function from dve.core_engine.backends.implementations.duckdb.duckdb_helpers import ( + duckdb_check_entity_empty, duckdb_record_index, duckdb_write_parquet, get_duckdb_type_from_annotation, @@ -18,6 +19,7 @@ from dve.core_engine.type_hints import URI, EntityName +@duckdb_check_entity_empty @duckdb_record_index @duckdb_write_parquet class DuckDBJSONReader(BaseFileReader): diff --git a/src/dve/core_engine/backends/implementations/duckdb/readers/xml.py b/src/dve/core_engine/backends/implementations/duckdb/readers/xml.py index 42e281a..ac11169 100644 --- a/src/dve/core_engine/backends/implementations/duckdb/readers/xml.py +++ b/src/dve/core_engine/backends/implementations/duckdb/readers/xml.py @@ -10,7 +10,10 @@ from dve.core_engine.backends.base.reader import read_function from dve.core_engine.backends.exceptions import MessageBearingError -from dve.core_engine.backends.implementations.duckdb.duckdb_helpers import duckdb_write_parquet +from dve.core_engine.backends.implementations.duckdb.duckdb_helpers import ( + duckdb_check_entity_empty, + duckdb_write_parquet, +) from dve.core_engine.backends.readers.xml import XMLStreamReader from dve.core_engine.backends.utilities import ( get_polars_type_from_annotation, @@ -20,6 +23,7 @@ from dve.core_engine.type_hints import URI +@duckdb_check_entity_empty @polars_record_index @duckdb_write_parquet class DuckDBXMLStreamReader(XMLStreamReader): diff --git a/src/dve/core_engine/backends/implementations/spark/readers/csv.py b/src/dve/core_engine/backends/implementations/spark/readers/csv.py index 2df30c5..5cd2f56 100644 --- a/src/dve/core_engine/backends/implementations/spark/readers/csv.py +++ b/src/dve/core_engine/backends/implementations/spark/readers/csv.py @@ -12,6 +12,7 @@ from dve.core_engine.backends.exceptions import EmptyFileError from dve.core_engine.backends.implementations.spark.spark_helpers import ( get_type_from_annotation, + spark_check_entity_empty, spark_record_index, spark_write_parquet, ) @@ -20,6 +21,7 @@ from dve.parser.file_handling import get_content_length +@spark_check_entity_empty @spark_record_index @spark_write_parquet class SparkCSVReader(CSVFileReader): diff --git a/src/dve/core_engine/backends/implementations/spark/readers/json.py b/src/dve/core_engine/backends/implementations/spark/readers/json.py index 6123009..6404231 100644 --- a/src/dve/core_engine/backends/implementations/spark/readers/json.py +++ b/src/dve/core_engine/backends/implementations/spark/readers/json.py @@ -11,6 +11,7 @@ from dve.core_engine.backends.exceptions import EmptyFileError from dve.core_engine.backends.implementations.spark.spark_helpers import ( get_type_from_annotation, + spark_check_entity_empty, spark_record_index, spark_write_parquet, ) @@ -18,6 +19,7 @@ from dve.parser.file_handling import get_content_length +@spark_check_entity_empty @spark_record_index @spark_write_parquet class SparkJSONReader(BaseFileReader): diff --git a/src/dve/core_engine/backends/implementations/spark/readers/xml.py b/src/dve/core_engine/backends/implementations/spark/readers/xml.py index ba42d29..4d6df6a 100644 --- a/src/dve/core_engine/backends/implementations/spark/readers/xml.py +++ b/src/dve/core_engine/backends/implementations/spark/readers/xml.py @@ -17,6 +17,7 @@ from dve.core_engine.backends.implementations.spark.spark_helpers import ( df_is_empty, get_type_from_annotation, + spark_check_entity_empty, spark_record_index, spark_write_parquet, ) @@ -29,6 +30,7 @@ """The mode to use when parsing XML files with Spark.""" +@spark_check_entity_empty @spark_record_index @spark_write_parquet class SparkXMLStreamReader(XMLStreamReader): @@ -55,6 +57,7 @@ def read_to_dataframe( ) +@spark_check_entity_empty @spark_record_index @spark_write_parquet class SparkXMLReader(BasicXMLFileReader): # pylint: disable=too-many-instance-attributes diff --git a/src/dve/core_engine/backends/implementations/spark/spark_helpers.py b/src/dve/core_engine/backends/implementations/spark/spark_helpers.py index 8c14132..50a6df4 100644 --- a/src/dve/core_engine/backends/implementations/spark/spark_helpers.py +++ b/src/dve/core_engine/backends/implementations/spark/spark_helpers.py @@ -456,6 +456,16 @@ def spark_get_entity_count(cls): return cls +def _spark_check_entity_empty(self, entity: DataFrame) -> bool: # pylint: disable=W0613 + return entity.count() == 0 + + +def spark_check_entity_empty(cls): + """Class decorator to check whether a supplied entity is empty""" + cls.check_entity_empty = _spark_check_entity_empty + return cls + + def get_all_registered_udfs(spark: SparkSession) -> set[str]: """Function to supply the names of a registered functions stored in the supplied spark session. diff --git a/src/dve/core_engine/configuration/v1/__init__.py b/src/dve/core_engine/configuration/v1/__init__.py index 10a452d..b743c27 100644 --- a/src/dve/core_engine/configuration/v1/__init__.py +++ b/src/dve/core_engine/configuration/v1/__init__.py @@ -57,6 +57,8 @@ """The operation """ RuleType = type[AbstractStep] """The metadata step type implemented by the rule.""" +AllowedAdditionalChecks = Literal["check_empty"] +"""Additional checks to be performed in the file_transformation stage""" class _BaseTypeDefintion(BaseModel): @@ -146,6 +148,11 @@ class _SchemaConfig(BaseModel): """A list of the field names within the schema which _must_ be provided.""" +class _ReaderAdditionalChecksConfig(BaseModel): + error_code: str + error_message: str + + class _ReaderConfig(BaseModel): # type: ignore """Reader configuration options for a model.""" @@ -155,6 +162,10 @@ class _ReaderConfig(BaseModel): # type: ignore """Keyword arguments for the reader.""" field_names: Optional[list[str]] = None """The field names to request from the reader. These are deprecated and will not be used.""" + additional_checks: dict[AllowedAdditionalChecks, _ReaderAdditionalChecksConfig] = Field( + default_factory=dict + ) + """Additional checks to be performed after the entity is read""" class _ModelConfig(_SchemaConfig): diff --git a/src/dve/pipeline/pipeline.py b/src/dve/pipeline/pipeline.py index ee5a6bc..7e8b8ae 100644 --- a/src/dve/pipeline/pipeline.py +++ b/src/dve/pipeline/pipeline.py @@ -238,11 +238,14 @@ def write_file_to_parquet( model_name, stringify_model(model), # type: ignore get_all_model_fields(models.values()), # type: ignore + dataset[model_name].reader_config[f".{ext.lower()}"].additional_checks, ), f"{out}{model_name}", ) except MessageBearingError as exc: - self._logger.error(f"Unable to process {model_name}", exc_info=exc) + self._logger.error( + f"While processing {model_name}, an issue was encountered", exc_info=exc + ) errors.extend(exc.messages) return list(dict.fromkeys(errors)) # remove any duplicate errors diff --git a/tests/features/flights.feature b/tests/features/flights.feature index b3e6790..2476445 100644 --- a/tests/features/flights.feature +++ b/tests/features/flights.feature @@ -263,3 +263,17 @@ Feature: Pipeline tests using the flights dataset | number_submission_rejections | 0 | | number_record_rejections | 14 | | number_warnings | 0 | + + Scenario: A flights submission where mandatory entity has no records submitted + Given I submit the flights file only_country_id.xml for processing + And A duckdb pipeline is configured with schema file 'flights_add_reader_checks.dischema.json' + And I add initial audit entries for the submission + Then the latest audit record for the submission is marked with processing status file_transformation + When I run the file transformation phase + Then the country entity is stored as a parquet after the file_transformation phase + And there are errors with the following details and associated error_count from the file_transformation phase + | FailureType | ErrorCode | error_count | + | submission | AIRPORTEMPTY | 1 | + And the latest audit record for the submission is marked with processing status error_report + When I run the error report phase + Then An error report is produced diff --git a/tests/features/steps/steps_pipeline.py b/tests/features/steps/steps_pipeline.py index c71bd45..8f40416 100644 --- a/tests/features/steps/steps_pipeline.py +++ b/tests/features/steps/steps_pipeline.py @@ -33,8 +33,8 @@ from utilities import ( load_errors_from_service, - get_test_file_path, SERVICE_TO_STORAGE_PATH_MAPPING, + get_test_file_path, get_all_errors_df, ) diff --git a/tests/features/steps/utilities.py b/tests/features/steps/utilities.py index 58edc67..dccefda 100644 --- a/tests/features/steps/utilities.py +++ b/tests/features/steps/utilities.py @@ -49,7 +49,7 @@ def load_errors_from_service(processing_folder: Path, service: str) -> pl.DataFr err_location = Path( processing_folder, "errors", - f"{SERVICE_TO_STORAGE_PATH_MAPPING.get(service, service)}_errors.jsonl", + f"{service}_errors.jsonl", ) msgs = [] try: diff --git a/tests/testdata/flights/flights_add_reader_checks.dischema.json b/tests/testdata/flights/flights_add_reader_checks.dischema.json new file mode 100644 index 0000000..8b4d1d0 --- /dev/null +++ b/tests/testdata/flights/flights_add_reader_checks.dischema.json @@ -0,0 +1,208 @@ +{ + "contract": { + "schemas": { + "passengers": { + "fields": { + "flight_id": "int", + "passenger_id": "int", + "passenger_name": "str" + } + } + }, + "error_details": "flights_data_contract_error_details.json", + "datasets": { + "country": { + "fields": { + "country_id": "int", + "country_name": "str" + }, + "reader_config": { + ".xml": { + "reader": "DuckDBXMLStreamReader", + "kwargs": { + "record_tag": "country", + "root_tag": "country" + } + } + }, + "key_field": "country_id", + "mandatory_fields": [ + "country_id", + "country_name" + ] + }, + "airport": { + "fields": { + "country_id": "int", + "airport_id": "int", + "airport_name": "str", + "postcode": "str" + }, + "reader_config": { + ".xml": { + "reader": "DuckDBXMLStreamReader", + "kwargs": { + "record_tag": "airport", + "root_tag": "country" + }, + "additional_checks": { + "check_empty": { + "error_code": "AIRPORTEMPTY", + "error_message": "No airport records included in submission" + } + } + } + }, + "key_field": "airport_id", + "mandatory_fields": [ + "airport_id" + ] + }, + "staff": { + "fields": { + "airport_id": "int", + "staff_id": "int", + "staff_name": "str", + "role": "str" + }, + "reader_config": { + ".xml": { + "reader": "DuckDBXMLStreamReader", + "kwargs": { + "record_tag": "staff_member", + "root_tag": "country" + } + } + }, + "key_field": "staff_id" + }, + "flights": { + "fields": { + "airport_id": "int", + "flight_id": "int", + "destination": "str" + }, + "reader_config": { + ".xml": { + "reader": "DuckDBXMLStreamReader", + "kwargs": { + "record_tag": "flight", + "root_tag": "country" + } + } + }, + "key_field": "flight_id" + }, + "passengers": { + "fields": { + "flight_id": "int", + "passenger_id": "int", + "passenger_name": "str" + }, + "reader_config": { + ".xml": { + "reader": "DuckDBXMLStreamReader", + "kwargs": { + "record_tag": "passenger", + "root_tag": "country" + } + } + }, + "key_field": "passenger_id" + } + } + }, + "transformations": { + "parameters": { + "entity": "country" + }, + "filters": [ + { + "entity": "flights", + "name": "flight_missing_id", + "expression": "flight_id IS NOT NULL", + "failure_type": "record", + "failure_message": "Record Rejected - Flight is missing an id", + "reporting_field": "flight_id", + "reporting_entity": "flights", + "category": "Blank", + "error_code": "FlightIDMissing" + }, + { + "entity": "flights", + "name": "invalid_destination", + "expression": "lower(destination) IN ('paris', 'madrid', 'new york', 'amsterdam', 'rome', 'dubai', 'dublin', 'lisbon', 'toronto')", + "failure_type": "record", + "failure_message": "Record Rejected - {{ destination }} is not a valid destination", + "reporting_field": "flight_id", + "reporting_entity": "flights", + "category": "Bad value", + "error_code": "InvalidFlightDestination" + }, + { + "entity": "passengers", + "name": "passenger_name_is_null", + "expression": "passenger_name IS NOT NULL", + "failure_type": "record", + "failure_message": "Record Rejected - Passenger Name is missing", + "reporting_field": "passenger_name", + "reporting_entity": "passengers", + "category": "Blank", + "error_code": "PassengerNameMissing" + }, + { + "entity": "staff", + "name": "staff_id_is_null", + "expression": "staff_id IS NOT NULL", + "failure_type": "record", + "failure_message": "Record Rejected - staff_id is missing", + "reporting_field": "passenger_name", + "reporting_entity": "passengers", + "category": "Blank", + "error_code": "StaffIDMissing" + } + ] + }, + "entity_relationships": { + "airport": { + "parent_entity": "country", + "join_fields": { + "country_id": "country_id" + }, + "mandatory": true, + "missing_parent_id_error_code": "AirportHasNoCountry", + "missing_parent_id_error_message": "Record rejected - No valid country id found for airport", + "no_valid_records_error_code": "CountryHasNoAirport", + "no_valid_records_error_message": "Group rejected - Unable to find any valid airports" + }, + "staff": { + "parent_entity": "airport", + "join_fields": { + "airport_id": "airport_id" + }, + "mandatory": true, + "missing_parent_id_error_code": "StaffHasNoAirport", + "missing_parent_id_error_message": "Record rejected - No valid airport id found for staff. Airport ID = {{ airport_id }}, Staff ID = {{ staff_id }}", + "no_valid_records_error_code": "AirportHasNoStaff", + "no_valid_records_error_message": "Group rejected - Airport has no valid staff. Airport ID = {{ airport_id }}" + }, + "flights": { + "parent_entity": "airport", + "join_fields": { + "airport_id": "airport_id" + }, + "mandatory": false, + "missing_parent_id_error_code": "FlightHasNoAirport", + "missing_parent_id_error_message": "Record Rejected - No valid airport found for flight" + }, + "passengers": { + "parent_entity": "flights", + "join_fields": { + "flight_id": "flight_id" + }, + "mandatory": false, + "missing_parent_id_error_code": "PassengerHasNoFlight", + "missing_parent_id_error_message": "Record rejected - No valid flight found for passenger" + } + } +} \ No newline at end of file From bf1c6c490cfa4eb92077fd75c04ea4ce4c644f1f Mon Sep 17 00:00:00 2001 From: stevenhsd <56357022+stevenhsd@users.noreply.github.com> Date: Tue, 29 Sep 2026 17:13:46 +0100 Subject: [PATCH 2/2] feat: following review, tweak location of additional reader checks config and add json schema support --- .../components/base_entity.schema.json | 9 +++++++- .../reader_additional_checks.schema.json | 22 +++++++++++++++++++ src/dve/core_engine/backends/base/reader.py | 7 ++++-- .../core_engine/configuration/v1/__init__.py | 10 ++++----- src/dve/pipeline/pipeline.py | 2 +- .../flights_add_reader_checks.dischema.json | 12 +++++----- 6 files changed, 47 insertions(+), 15 deletions(-) create mode 100644 docs/advanced_guidance/json_schemas/contract/components/reader_additional_checks.schema.json diff --git a/docs/advanced_guidance/json_schemas/contract/components/base_entity.schema.json b/docs/advanced_guidance/json_schemas/contract/components/base_entity.schema.json index b89291a..06607cb 100644 --- a/docs/advanced_guidance/json_schemas/contract/components/base_entity.schema.json +++ b/docs/advanced_guidance/json_schemas/contract/components/base_entity.schema.json @@ -26,8 +26,15 @@ "items": { "type": "string" } + }, + "reader_additional_checks": { + "description": "A mapping of additional checks to perform on entities after initial read", + "type": "object", + "additionalProperties": { + "$ref": "reader_additional_checks.json" } - }, + } +}, "required": ["fields"] } diff --git a/docs/advanced_guidance/json_schemas/contract/components/reader_additional_checks.schema.json b/docs/advanced_guidance/json_schemas/contract/components/reader_additional_checks.schema.json new file mode 100644 index 0000000..6c01eb7 --- /dev/null +++ b/docs/advanced_guidance/json_schemas/contract/components/reader_additional_checks.schema.json @@ -0,0 +1,22 @@ +{ + "$schema": "https://json-schema.org/draft/2020-12/schema", + "$id": "data-ingest:contract/components/reader_additional_checks.schema.json", + "title": "reader_additional_checks", + "description": "Additional checks to perform on initially read entities", + "type": "object", + "properties": { + "error_code": { + "description": "The code to be used for the additional check specified", + "type": "string" + }, + "error_message": { + "description": "The message to be displayed for the additional check specified.", + "type": "string" + } + }, + "required": [ + "error_code", + "error_message" + ], + "additionalProperties": false +} \ No newline at end of file diff --git a/src/dve/core_engine/backends/base/reader.py b/src/dve/core_engine/backends/base/reader.py index e65fb76..ff3586c 100644 --- a/src/dve/core_engine/backends/base/reader.py +++ b/src/dve/core_engine/backends/base/reader.py @@ -10,7 +10,10 @@ from dve.core_engine.backends.exceptions import MessageBearingError, ReaderLacksEntityTypeSupport from dve.core_engine.backends.types import EntityName, EntityType -from dve.core_engine.configuration.v1 import AllowedAdditionalChecks, _ReaderAdditionalChecksConfig +from dve.core_engine.configuration.v1 import ( + AllowedAdditionalReaderChecks, + _ReaderAdditionalChecksConfig, +) from dve.core_engine.message import FeedbackMessage from dve.core_engine.type_hints import URI, ArbitraryFunction, WrapDecorator from dve.parser.file_handling.service import open_stream @@ -111,7 +114,7 @@ def read_to_entity_type( schema: type[BaseModel], all_model_fields: Optional[set[str]] = None, additional_checks: Optional[ - dict[AllowedAdditionalChecks, _ReaderAdditionalChecksConfig] + dict[AllowedAdditionalReaderChecks, _ReaderAdditionalChecksConfig] ] = None, ): """Read to the specified entity type, if supported. diff --git a/src/dve/core_engine/configuration/v1/__init__.py b/src/dve/core_engine/configuration/v1/__init__.py index b743c27..9a400b5 100644 --- a/src/dve/core_engine/configuration/v1/__init__.py +++ b/src/dve/core_engine/configuration/v1/__init__.py @@ -57,7 +57,7 @@ """The operation """ RuleType = type[AbstractStep] """The metadata step type implemented by the rule.""" -AllowedAdditionalChecks = Literal["check_empty"] +AllowedAdditionalReaderChecks = Literal["check_empty"] """Additional checks to be performed in the file_transformation stage""" @@ -162,10 +162,6 @@ class _ReaderConfig(BaseModel): # type: ignore """Keyword arguments for the reader.""" field_names: Optional[list[str]] = None """The field names to request from the reader. These are deprecated and will not be used.""" - additional_checks: dict[AllowedAdditionalChecks, _ReaderAdditionalChecksConfig] = Field( - default_factory=dict - ) - """Additional checks to be performed after the entity is read""" class _ModelConfig(_SchemaConfig): @@ -177,6 +173,10 @@ class _ModelConfig(_SchemaConfig): """A single key field to be used by the model.""" reader_config: dict[Extension, _ReaderConfig] """Reader configuration options for the model.""" + reader_additional_checks: dict[AllowedAdditionalReaderChecks, _ReaderAdditionalChecksConfig] = ( + Field(default_factory=dict) + ) + """Additional checks to be performed after the entity is read""" aliases: dict[FieldName, FieldName] = Field(default_factory=dict) """An alias field name mapping.""" diff --git a/src/dve/pipeline/pipeline.py b/src/dve/pipeline/pipeline.py index 7e8b8ae..455ed8d 100644 --- a/src/dve/pipeline/pipeline.py +++ b/src/dve/pipeline/pipeline.py @@ -238,7 +238,7 @@ def write_file_to_parquet( model_name, stringify_model(model), # type: ignore get_all_model_fields(models.values()), # type: ignore - dataset[model_name].reader_config[f".{ext.lower()}"].additional_checks, + dataset[model_name].reader_additional_checks, ), f"{out}{model_name}", ) diff --git a/tests/testdata/flights/flights_add_reader_checks.dischema.json b/tests/testdata/flights/flights_add_reader_checks.dischema.json index 8b4d1d0..5f8e63b 100644 --- a/tests/testdata/flights/flights_add_reader_checks.dischema.json +++ b/tests/testdata/flights/flights_add_reader_checks.dischema.json @@ -44,15 +44,15 @@ "kwargs": { "record_tag": "airport", "root_tag": "country" - }, - "additional_checks": { - "check_empty": { - "error_code": "AIRPORTEMPTY", - "error_message": "No airport records included in submission" - } } } }, + "reader_additional_checks": { + "check_empty": { + "error_code": "AIRPORTEMPTY", + "error_message": "No airport records included in submission" + } + }, "key_field": "airport_id", "mandatory_fields": [ "airport_id"