Improvement/ add pandera schema to rapid dataframe schema

.gitignore

@@ -186,3 +186,5 @@ docs/_build/
 frontend/playwright/.auth
 frontend/playwright/.downloads
 frontend/test-results/
+
+backend/api/pandera_checks

backend/api/adapter/dynamodb_adapter.py

@@ -85,7 +85,9 @@ def update_job(self, job: Job) -> None:
         pass
 
     @abstractmethod
-    def get_latest_successful_upload_job(self, dataset: Type[DatasetMetadata]) -> Optional[Dict]:
+    def get_latest_successful_upload_job(
+        self, dataset: Type[DatasetMetadata]
+    ) -> Optional[Dict]:
         pass
 
     @abstractmethod
@@ -177,14 +179,15 @@ def store_schema(self, schema: Schema) -> None:
             AppLogger.info(
                 f"Storing schema for {schema.metadata.string_representation()}"
             )
-            self.schema_table.put_item(
-                Item={
-                    "PK": schema.metadata.dataset_identifier(with_version=False),
-                    "SK": schema.metadata.get_version(),
-                    **schema.metadata.model_dump(),
-                    COLUMNS: [col.model_dump() for col in schema.columns],
-                }
-            )
+            item = {
+                "PK": schema.metadata.dataset_identifier(with_version=False),
+                "SK": schema.metadata.get_version(),
+                **schema.metadata.model_dump(),
+                COLUMNS: [col.model_dump() for col in schema.columns],
+            }
+            if schema.panderaDataFrameSchema is not None:
+                item["panderaDataFrameSchema"] = schema.panderaDataFrameSchema.to_json()
+            self.schema_table.put_item(Item=item)
         except ClientError as error:
             self._handle_client_error(
                 f"Error storing schema for {schema.metadata.string_representation()}",
@@ -389,7 +392,9 @@ def get_job(self, job_id: str) -> Dict:
         except ClientError as error:
             self._handle_client_error("Error fetching job from the database", error)
 
-    def get_latest_successful_upload_job(self, dataset: Type[DatasetMetadata]) -> Optional[Dict]:
+    def get_latest_successful_upload_job(
+        self, dataset: Type[DatasetMetadata]
+    ) -> Optional[Dict]:
         """
         Get the most recent successful upload job for a specific dataset.
         Returns the job details including subject_id (uploader) or None if no successful upload exists.
@@ -404,13 +409,15 @@ def get_latest_successful_upload_job(self, dataset: Type[DatasetMetadata]) -> Op
                     & Attr("Layer").eq(dataset.layer)
                     & Attr("Domain").eq(dataset.domain)
                     & Attr("Dataset").eq(dataset.dataset)
-                )
+                ),
             )
 
             if not jobs:
                 return None
 
-            sorted_jobs = sorted(jobs, key=lambda x: x.get("CreatedAt", 0), reverse=True)
+            sorted_jobs = sorted(
+                jobs, key=lambda x: x.get("CreatedAt", 0), reverse=True
+            )
             return self._map_job(sorted_jobs[0])
         except ClientError as error:
             AppLogger.warning(f"Error fetching latest upload job for dataset: {error}")

backend/api/application/services/dataset_validation.py

@@ -1,6 +1,6 @@
 import re
 from typing import Tuple
-
+import json
 import pandas as pd
 from pandas import Timestamp
 import pandera
@@ -18,6 +18,7 @@
 )
 from api.domain.schema import Schema
 from api.domain.validation_context import ValidationContext
+from api.common.logger import AppLogger
 
 
 def build_validated_dataframe(schema: Schema, dataframe: pd.DataFrame) -> pd.DataFrame:
@@ -35,6 +36,7 @@ def transform_and_validate(schema: Schema, data: pd.DataFrame) -> pd.DataFrame:
         .pipe(dataset_has_correct_data_types, schema)
         .pipe(dataset_has_no_illegal_characters_in_partition_columns, schema)
         .pipe(validate_with_pandera, schema)
+        .pipe(validate_with_pandera_schema, schema)
     )
 
     if validation_context.has_errors():
@@ -120,7 +122,9 @@ def dataset_has_no_illegal_characters_in_partition_columns(
     error_list = []
     for column in schema.get_partition_columns():
         series = data_frame[column.name]
-        if not column.is_of_data_type(DateType) and pd.api.types.is_string_dtype(series):
+        if not column.is_of_data_type(DateType) and pd.api.types.is_string_dtype(
+            series
+        ):
             any_illegal_characters = series.str.contains("/", na=False).any()
             if any_illegal_characters:
                 error_list.append(
@@ -183,7 +187,7 @@ def parse_pandera_errors(exc: pandera.errors.SchemaErrors) -> list[str]:
 
     # Creating a list of singular (json like) entries from the pandera error string
     # For example: {'check': 'pandera_check', 'error': 'error message', ...}
-    failure_object_pattern = r'\{\s*(?:[^{}]*?)\}'
+    failure_object_pattern = r"\{\s*(?:[^{}]*?)\}"
     failure_objects = re.findall(failure_object_pattern, error_str)
 
     # Extracting and cleaning each error statement
@@ -195,15 +199,15 @@ def parse_pandera_errors(exc: pandera.errors.SchemaErrors) -> list[str]:
             continue
 
         error_msg = error_match.group(1)
-        error_msg = error_msg.replace(r"\'", "'").replace(r'\"', '"')
+        error_msg = error_msg.replace(r"\'", "'").replace(r"\"", '"')
 
         check_match = re.search(r'"check":\s*"([^"]*)"', obj)
         check_name = check_match.group(1) if check_match else None
 
-        if ':' in error_msg and ('Name:' in error_msg or 'dtype:' in error_msg):
-            error_msg = error_msg.split(':')[0]
+        if ":" in error_msg and ("Name:" in error_msg or "dtype:" in error_msg):
+            error_msg = error_msg.split(":")[0]
 
-        if check_name and check_name not in ['not_nullable', 'field_uniqueness']:
+        if check_name and check_name not in ["not_nullable", "field_uniqueness"]:
             error_msg = f"[{check_name}] {error_msg}"
 
         error_messages.append(error_msg)
@@ -221,3 +225,37 @@ def validate_with_pandera(
     except pandera.errors.SchemaErrors as exc:
         error_list = parse_pandera_errors(exc)
         return data_frame, error_list
+
+
+def parse_pandera_schema_errors(exc: pandera.errors.SchemaErrors) -> list[str]:
+
+    error_messages = []
+
+    for error_reason in exc.message["DATA"]:
+        if error_reason == "CHECK_ERROR":
+            for error_entry in exc.message["DATA"][error_reason]:
+                if error_entry["column"] is not None:
+                    error_messages.append(
+                        f"Column '{error_entry['column']}' internal error in check: {error_entry['check']}"
+                    )
+                else:
+                    error_messages.append(
+                        f"Internal error in check: {error_entry['check']}"
+                    )
+        else:
+            for error_entry in exc.message["DATA"][error_reason]:
+                error_messages.append(error_entry["error"])
+
+    return error_messages
+
+
+def validate_with_pandera_schema(
+    data_frame: pd.DataFrame, schema: Schema
+) -> Tuple[pd.DataFrame, list[str]]:
+    try:
+        validated_df = schema.pandera_schema_validate(data_frame, lazy=True)
+        return validated_df, []
+    except pandera.errors.SchemaErrors as exc:
+        AppLogger.info(f"    Pandera schema entire error: {exc.message}")
+        error_list = parse_pandera_schema_errors(exc)
+        return data_frame, error_list

backend/api/application/services/schema_service.py

@@ -56,6 +56,7 @@ def _parse_schema(self, schema: dict, only_metadata: bool = False):
         return Schema(
             metadata=metadata,
             columns=[Column.model_validate(col) for col in schema[COLUMNS]],
+            panderaDataFrameSchema=schema.get("panderaDataFrameSchema", None),
         )
 
     def get_schema_metadatas(

backend/api/domain/schema.py

@@ -3,8 +3,11 @@
 
 import awswrangler as wr
 from pydantic.main import BaseModel
+from pydantic import field_serializer, field_validator
 import pyarrow as pa
 import pandera
+import pandera.pandas as pandera_pandas
+import pandera.io.pandas_io as pandera_io
 
 from api.domain.schema_metadata import Owner, SchemaMetadata
 from rapid.items.schema import Column, UpdateBehaviour
@@ -16,6 +19,21 @@
 class Schema(BaseModel):
     metadata: SchemaMetadata
     columns: List[Column]
+    panderaDataFrameSchema: Optional[pandera_pandas.DataFrameSchema] = None
+
+    @field_serializer("panderaDataFrameSchema", mode="plain")
+    def pandera_dump(self, value: pandera_pandas.DataFrameSchema) -> str:
+        if value is not None:
+            return value.to_json()
+        else:
+            return value
+
+    @field_validator("panderaDataFrameSchema", mode="before")
+    def pandera_load(cls, value: str) -> pandera_pandas.DataFrameSchema:
+        if type(value) is str:
+            return pandera_io.from_json(value)
+        else:
+            return value
 
     def get_layer(self) -> str:
         return self.metadata.get_layer()
@@ -100,9 +118,14 @@ def generate_storage_schema(self) -> pa.schema:
         )
 
     def pandera_validate(self, df, **kwargs):
-        pandera_columns = {
-            col.name: col.to_pandera_column()
-            for col in self.columns
-        }
-        pandera_schema = pandera.DataFrameSchema(metadata=self.metadata, columns=pandera_columns)
+        pandera_columns = {col.name: col.to_pandera_column() for col in self.columns}
+        pandera_schema = pandera.DataFrameSchema(
+            metadata=self.metadata, columns=pandera_columns
+        )
         return pandera_schema.validate(df, **kwargs)
+
+    def pandera_schema_validate(self, df, **kwargs):
+        if self.panderaDataFrameSchema is not None:
+            return self.panderaDataFrameSchema.validate(df, **kwargs)
+        else:
+            return df

backend/api/entry.py

@@ -42,6 +42,7 @@
 from api.controller.subjects import subjects_router
 from api.controller.user import user_router
 from api.exception_handler import add_exception_handlers
+from api.pandera_custom_checks_load import pandera_custom_checks_load
 
 try:
     load_dotenv()
@@ -58,9 +59,7 @@
 permissions_service = PermissionsService()
 upload_service = DatasetAccessEvaluator()
 
-app = FastAPI(
-    openapi_url=f"{BASE_API_PATH}/openapi.json", docs_url=None
-)
+app = FastAPI(openapi_url=f"{BASE_API_PATH}/openapi.json", docs_url=None)
 app.mount("/static", StaticFiles(directory="static"), name="static")
 app.openapi = custom_openapi_docs_generator(app)
 add_exception_handlers(app)
@@ -79,6 +78,10 @@
 @app.on_event("startup")
 async def startup_event():
     init_logger()
+    try:
+        pandera_custom_checks_load()
+    except Exception as e:
+        AppLogger.info(f"Failed to load pandera custom checks due to: {str(e)}")
 
 
 @app.middleware("http")
@@ -247,9 +250,9 @@ def _set_security_headers(response) -> None:
         "img-src 'self' data: "
         "fastapi.tiangolo.com/img/favicon.png;"
     )
-    response.headers[
-        "Strict-Transport-Security"
-    ] = "max-age=31536000 ; includeSubDomains"
+    response.headers["Strict-Transport-Security"] = (
+        "max-age=31536000 ; includeSubDomains"
+    )
     response.headers["X-Content-Type-Options"] = "nosniff"
     response.headers["X-Permitted-Cross-Domain-Policies"] = "none"
     response.headers["Referrer-Policy"] = "strict-origin"

backend/api/pandera_custom_checks_load.py

@@ -0,0 +1,21 @@
+import os
+import json
+import importlib
+
+
+def pandera_custom_checks_load():
+    pandera_files = os.getenv("PANDERA_FILES", None)
+    if pandera_files is None:
+        return
+
+    pandera_files = json.loads(pandera_files)
+
+    for file in pandera_files:
+        file_path = os.path.join("api", "pandera_checks", file)
+
+        os.makedirs(os.path.dirname(file_path), exist_ok=True)
+
+        with open(file_path, "w") as f:
+            f.write(pandera_files[file])
+
+        importlib.import_module(file_path.replace("/", ".").replace("\\", ".")[:-3])

backend/rapid/items/schema.py

@@ -1,8 +1,10 @@
 # Note: This class is replicated in the api code, they should be de-duplicated once the external dependencies are removed from the API
 from strenum import StrEnum
 from typing import Dict, List, Optional, Union, Any
-from pydantic import BaseModel, ConfigDict
+from pydantic import BaseModel, ConfigDict, field_serializer, field_validator
 import pandera
+import pandera.pandas as pandera_pandas
+import pandera.io.pandas_io as pandera_io
 
 
 class SensitivityLevel(StrEnum):
@@ -48,7 +50,7 @@ class Column(BaseModel):
 
     def is_of_data_type(self, d_type: StrEnum) -> bool:
         return self.data_type in list(d_type)
-    
+
     def to_pandera_column(self) -> pandera.Column:
         """
         Convert Column to Pandera Column for Pandera data validation.
@@ -96,8 +98,10 @@ def _dict_to_pandera_check(self, check_dict: Dict[str, Any]) -> pandera.Check:
             pattern = params.get("pattern")
             return pandera.Check.str_matches(pattern)
         else:
-            raise ValueError(f"Unsupported check type: {check_type}. Valid types are: "
-                             "in_range, isin, str_length, greater_than, less_than, str_matches.")
+            raise ValueError(
+                f"Unsupported check type: {check_type}. Valid types are: "
+                "in_range, isin, str_length, greater_than, less_than, str_matches."
+            )
 
 
 class Schema(BaseModel):
@@ -143,6 +147,21 @@ class Schema(BaseModel):
 
     metadata: SchemaMetadata
     columns: List[Column]
+    panderaDataFrameSchema: Optional[pandera_pandas.DataFrameSchema] = None
+
+    @field_serializer("panderaDataFrameSchema", mode="plain")
+    def pandera_dump(self, value: pandera_pandas.DataFrameSchema) -> str:
+        if value is not None:
+            return value.to_json()
+        else:
+            return value
+
+    @field_validator("panderaDataFrameSchema", mode="before")
+    def pandera_load(cls, value: str) -> pandera_pandas.DataFrameSchema:
+        if type(value) is str:
+            return pandera_io.from_json(value)
+        else:
+            return value
 
     def are_columns_the_same(
         self, new_columns: Union[List[Column], List[dict]]
@@ -179,9 +198,8 @@ def pandera_validate(self, df, **kwargs):
         Raises:
             pandera.errors.SchemaErrors: If validation fails
         """
-        pandera_columns = {
-            col.name: col.to_pandera_column()
-            for col in self.columns
-        }
-        pandera_schema = pandera.DataFrameSchema(metadata=self.metadata, columns=pandera_columns)
+        pandera_columns = {col.name: col.to_pandera_column() for col in self.columns}
+        pandera_schema = pandera.DataFrameSchema(
+            metadata=self.metadata, columns=pandera_columns
+        )
         return pandera_schema.validate(df, **kwargs)

backend/requirements-dev.txt

@@ -8,5 +8,6 @@ requests-mock
 setuptools
 twine
 pip-audit
+watchfiles
 
 -r ./requirements.txt

backend/requirements.txt

@@ -6,7 +6,7 @@ gunicorn
 httpx
 jinja2
 pandas
-pandera
+pandera[io]
 psutil
 pyarrow
 pyjwt
@@ -17,4 +17,3 @@ uvicorn
 requests
 strenum
 pytest-order
-