# File: pipeline_steps/config_cradle_data_load.py
from typing import List, Optional, Dict, Any, Set
import re
from datetime import datetime
from pydantic import (
BaseModel,
Field,
field_validator,
model_validator,
ConfigDict,
PrivateAttr,
)
from ...core.base.config_base import BasePipelineConfig
[docs]
def get_flattened_fields(config_obj, prefix="") -> Dict[str, List[str]]:
"""
Recursively gather all fields from a config object and its nested objects,
flattening them into a single list with dot notation to indicate hierarchy.
Args:
config_obj: Configuration object to analyze
prefix: String prefix for nested fields (for recursive calls)
Returns:
Dict with keys 'essential', 'system', and 'derived' mapping to lists of field names
"""
# Get the fields at this level
if hasattr(config_obj, "categorize_fields"):
categories = config_obj.categorize_fields()
else:
# Initialize empty categories if the object doesn't support categorization
categories = {"essential": [], "system": [], "derived": []}
# Add the prefix to the field names
result = {
"essential": [f"{prefix}{field}" for field in categories["essential"]],
"system": [f"{prefix}{field}" for field in categories["system"]],
"derived": [f"{prefix}{field}" for field in categories["derived"]],
}
# Handle nested configuration objects
for field_name in categories["essential"] + categories["system"]:
field_value = getattr(config_obj, field_name)
# Skip None values
if field_value is None:
continue
# Handle list of configuration objects
if isinstance(field_value, list) and len(field_value) > 0:
# Check if items have categorize_fields method
if all(
hasattr(item, "categorize_fields")
for item in field_value
if item is not None
):
for i, item in enumerate(field_value):
if item is None:
continue
nested_result = get_flattened_fields(
item, f"{prefix}{field_name}[{i}]."
)
for cat, fields in nested_result.items():
result[cat].extend(fields)
# Handle single configuration objects
elif hasattr(field_value, "categorize_fields"):
nested_result = get_flattened_fields(field_value, f"{prefix}{field_name}.")
for cat, fields in nested_result.items():
result[cat].extend(fields)
return result
[docs]
class BaseCradleComponentConfig(BaseModel):
"""
Base class for Cradle configuration components with three-tier field classification support.
Implements common functionality for categorizing fields and supporting inheritance.
Fields are organized into three tiers:
1. Tier 1: Essential User Inputs - fields that users must explicitly provide
2. Tier 2: System Inputs with Defaults - fields with reasonable defaults that can be overridden
3. Tier 3: Derived Fields - fields calculated from other fields (private attributes with properties)
"""
# Model configuration
model_config = ConfigDict(
arbitrary_types_allowed=True,
validate_assignment=True,
extra="allow", # Changed from "forbid" to "allow" to fix circular reference handling
)
[docs]
def categorize_fields(self) -> Dict[str, List[str]]:
"""
Categorize all fields into three tiers:
1. Tier 1: Essential User Inputs - fields with no defaults (required)
2. Tier 2: System Inputs - fields with defaults (optional)
3. Tier 3: Derived Fields - properties that access private attributes
Returns:
Dict with keys 'essential', 'system', and 'derived' mapping to lists of field names
"""
# Initialize categories
categories = {
"essential": [], # Tier 1: Required, public
"system": [], # Tier 2: Optional (has default), public
"derived": [], # Tier 3: Public properties
}
# Get model fields
model_fields = self.__class__.model_fields
# Categorize public fields into essential (required) or system (with defaults)
for field_name, field_info in model_fields.items():
# Skip private fields
if field_name.startswith("_"):
continue
# Use is_required() to determine if a field is essential
if field_info.is_required():
categories["essential"].append(field_name)
else:
categories["system"].append(field_name)
# Find derived properties (public properties that aren't in model_fields)
for attr_name in dir(self):
if (
not attr_name.startswith("_")
and attr_name not in model_fields
and isinstance(getattr(type(self), attr_name, None), property)
):
categories["derived"].append(attr_name)
return categories
[docs]
def get_public_init_fields(self) -> Dict[str, Any]:
"""
Get fields suitable for initializing a child config.
Only includes fields that should be passed to child class constructors.
Returns:
Dict[str, Any]: Dictionary of field names to values for child initialization
"""
# Use categorize_fields to get essential and system fields
categories = self.categorize_fields()
# Initialize result dict
init_fields = {}
# Add all essential fields (Tier 1)
for field_name in categories["essential"]:
init_fields[field_name] = getattr(self, field_name)
# Add all system fields (Tier 2) that aren't None
for field_name in categories["system"]:
value = getattr(self, field_name)
if value is not None: # Only include non-None values
init_fields[field_name] = value
return init_fields
[docs]
class MdsDataSourceConfig(BaseCradleComponentConfig):
"""
Configuration for MDS data source with three-tier field classification.
Fields are organized into three tiers:
1. Tier 1: Essential User Inputs - fields that users must explicitly provide
2. Tier 2: System Inputs with Defaults - fields with reasonable defaults that can be overridden
3. Tier 3: Derived Fields - fields calculated from other fields (private attributes with properties)
"""
# ===== Essential User Inputs (Tier 1) =====
# These are fields that users must explicitly provide
service_name: str = Field(description="Name of the MDS service")
region: str = Field(description="Region code for MDS (e.g. 'NA', 'EU', 'FE')")
output_schema: List[Dict[str, Any]] = Field(
description="List of dictionaries describing each output column, "
"e.g. [{'field_name':'objectId','field_type':'STRING'}, …]"
)
# ===== System Inputs with Defaults (Tier 2) =====
# These are fields with reasonable defaults that users can override
org_id: int = Field(
default=0,
description="Organization ID (integer) for MDS. Default as 0 for regional MDS bucket.",
)
use_hourly_edx_data_set: bool = Field(
default=False, description="Whether to use the hourly EDX dataset flag in MDS"
)
# ===== Derived Fields (Tier 3) =====
# None currently for this class
[docs]
@field_validator("region")
@classmethod
def validate_region(cls, v: str) -> str:
valid = {"NA", "EU", "FE"}
if v not in valid:
raise ValueError(f"region must be one of {valid}, got '{v}'")
return v
[docs]
class EdxDataSourceConfig(BaseCradleComponentConfig):
"""
Configuration for EDX data source with three-tier field classification.
Supports two input modes:
1. Direct ARN input: Provide edx_arn directly
2. Component-based input: Provide edx_provider, edx_subject, edx_dataset, edx_manifest_key
Fields are organized into three tiers:
1. Tier 1: Essential User Inputs - fields that users must explicitly provide
2. Tier 2: System Inputs with Defaults - fields with reasonable defaults that can be overridden
3. Tier 3: Derived Fields - fields calculated from other fields (private attributes with properties)
"""
# ===== Essential User Inputs (Tier 1) =====
# These are fields that users must always explicitly provide
# (None currently - all fields are optional with defaults)
# ===== System Inputs with Defaults (Tier 2) =====
# Control field that determines requirement of component fields
schema_overrides: Optional[List[Dict[str, Any]]] = Field(
default=None,
description=(
"List of dicts overriding the EDX schema, e.g. "
"[{'field_name':'order_id','field_type':'STRING'}, …]. "
"If None, EDX will use the default schema."
),
)
edx_arn: Optional[str] = Field(
default=None,
description="Complete EDX manifest ARN. If provided, individual components are ignored.",
)
# Conditionally required fields (required only when edx_arn is None)
edx_provider: Optional[str] = Field(
default=None,
description="Provider portion of the EDX manifest ARN (required if edx_arn not provided)",
)
edx_subject: Optional[str] = Field(
default=None,
description="Subject portion of the EDX manifest ARN (required if edx_arn not provided)",
)
edx_dataset: Optional[str] = Field(
default=None,
description="Dataset portion of the EDX manifest ARN (required if edx_arn not provided)",
)
edx_manifest_key: Optional[str] = Field(
default=None,
description="Manifest key in format '[\"xxx\",...]' (required if edx_arn not provided)",
)
# ===== Derived Fields (Tier 3) =====
# These are fields calculated from other fields
_edx_manifest: Optional[str] = PrivateAttr(default=None)
@property
def edx_manifest(self) -> str:
"""Get EDX manifest ARN from direct input or built from components."""
if self._edx_manifest is None:
if self.edx_arn is not None:
# Mode 1: Direct ARN input
self._edx_manifest = self.edx_arn
else:
# Mode 2: Build from components (existing logic)
self._edx_manifest = (
f"arn:amazon:edx:iad::manifest/"
f"{self.edx_provider}/{self.edx_subject}/{self.edx_dataset}/{self.edx_manifest_key}"
)
return self._edx_manifest
[docs]
def categorize_fields(self) -> Dict[str, List[str]]:
"""Dynamic field categorization based on edx_arn presence."""
categories = {
"essential": [], # No always-required fields
"system": [
"edx_arn",
"schema_overrides",
], # Control field and optional schema overrides
"derived": ["edx_manifest"], # Computed property
}
# Component fields are system-level but conditionally required
component_fields = [
"edx_provider",
"edx_subject",
"edx_dataset",
"edx_manifest_key",
]
if self.edx_arn is None:
# When no ARN provided, components become essential
categories["essential"].extend(component_fields)
else:
# When ARN provided, components are just system fields
categories["system"].extend(component_fields)
return categories
[docs]
class AndesDataSourceConfig(BaseCradleComponentConfig):
"""
Configuration for Andes data source with three-tier field classification.
Fields are organized into three tiers:
1. Tier 1: Essential User Inputs - fields that users must explicitly provide
2. Tier 2: System Inputs with Defaults - fields with reasonable defaults that can be overridden
3. Tier 3: Derived Fields - fields calculated from other fields (private attributes with properties)
"""
# ===== Essential User Inputs (Tier 1) =====
# These are fields that users must explicitly provide
provider: str = Field(description="Andes provider ID (32-digit UUID or 'booker')")
table_name: str = Field(description="Name of the Andes table")
# ===== System Inputs with Defaults (Tier 2) =====
# These are fields with reasonable defaults that users can override
andes3_enabled: bool = Field(
default=True, description="Whether the table uses Andes 3.0 with latest version"
)
# ===== Derived Fields (Tier 3) =====
# None currently for this class
# Model configuration overrides
model_config = ConfigDict(
arbitrary_types_allowed=True,
validate_assignment=True,
extra="allow", # Changed from "forbid" to "allow" to fix circular reference handling
str_strip_whitespace=True,
)
[docs]
@field_validator("provider")
@classmethod
def validate_provider(cls, v: str) -> str:
"""
Validate that the provider is either:
1. A valid 32-character UUID
2. The special case 'booker'
"""
if v == "booker":
return v
uuid_pattern = re.compile(
r"^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$"
)
if not uuid_pattern.match(v.lower()):
raise ValueError(
"provider must be either 'booker' or a valid 32-digit UUID "
"(8-4-4-4-12 format). "
"Verify provider validity at: "
f"https://datacentral.a2z.com/hoot/providers/{v}"
)
return v
[docs]
@field_validator("table_name")
@classmethod
def validate_table_name(cls, v: str) -> str:
"""
Validate that the table name is not empty and follows valid format.
"""
if not v or not v.strip():
raise ValueError("table_name cannot be empty")
return v
def __str__(self) -> str:
"""String representation of the Andes config."""
return (
f"AndesDataSourceConfig(provider='{self.provider}', "
f"table_name='{self.table_name}', "
f"andes3_enabled={self.andes3_enabled})"
)
[docs]
class DataSourceConfig(BaseCradleComponentConfig):
"""
Configuration for data sources with three-tier field classification.
Corresponds to com.amazon.secureaisandboxproxyservice.models.datasource.DataSource
Fields are organized into three tiers:
1. Tier 1: Essential User Inputs - fields that users must explicitly provide
2. Tier 2: System Inputs with Defaults - fields with reasonable defaults that can be overridden
3. Tier 3: Derived Fields - fields calculated from other fields (private attributes with properties)
"""
# ===== Essential User Inputs (Tier 1) =====
# These are fields that users must explicitly provide
data_source_name: str = Field(
description="Logical name for this data source (e.g. 'RAW_MDS_NA' or 'TAGS')"
)
data_source_type: str = Field(description="One of 'MDS', 'EDX', or 'ANDES'")
# ===== System Inputs with Defaults (Tier 2) =====
# These are fields with reasonable defaults that users can override
mds_data_source_properties: Optional[MdsDataSourceConfig] = Field(
default=None, description="If data_source_type=='MDS', this must be provided"
)
edx_data_source_properties: Optional[EdxDataSourceConfig] = Field(
default=None, description="If data_source_type=='EDX', this must be provided"
)
andes_data_source_properties: Optional[AndesDataSourceConfig] = Field(
default=None, description="If data_source_type=='ANDES', this must be provided"
)
# ===== Derived Fields (Tier 3) =====
# None currently for this class
# Override model_config to set frozen=True for this class
model_config = ConfigDict(
arbitrary_types_allowed=True,
validate_assignment=True,
extra="allow", # Changed from "forbid" to "allow" to fix circular reference handling
frozen=True,
)
[docs]
@field_validator("data_source_type")
@classmethod
def validate_type(cls, v: str) -> str:
# data_source_type is compared case-sensitively against uppercase
# literals downstream (check_properties: t == "MDS"/"EDX"/"ANDES") and is
# a SAIS Cradle SDK enum. Accept case-insensitive input but normalize to
# the canonical uppercase value so the downstream comparisons stay correct.
allowed: Set[str] = {"MDS", "EDX", "ANDES"}
v_upper = v.upper()
if v_upper not in allowed:
raise ValueError(
f"data_source_type must be one of {sorted(allowed)} (case-insensitive), got '{v}'"
)
return v_upper
[docs]
@model_validator(mode="after")
@classmethod
def check_properties(cls, model: "DataSourceConfig") -> "DataSourceConfig":
"""
Ensure the appropriate properties are set based on data_source_type
and that only one set of properties is provided.
"""
t = model.data_source_type
# Check required properties are present
if t == "MDS" and model.mds_data_source_properties is None:
raise ValueError(
"mds_data_source_properties must be set when data_source_type=='MDS'"
)
if t == "EDX" and model.edx_data_source_properties is None:
raise ValueError(
"edx_data_source_properties must be set when data_source_type=='EDX'"
)
if t == "ANDES" and model.andes_data_source_properties is None:
raise ValueError(
"andes_data_source_properties must be set when data_source_type=='ANDES'"
)
# Ensure only one set of properties is provided
properties_count = sum(
1
for prop in [
model.mds_data_source_properties,
model.edx_data_source_properties,
model.andes_data_source_properties,
]
if prop is not None
)
if properties_count > 1:
raise ValueError(
"Only one of mds_data_source_properties, edx_data_source_properties, "
"or andes_data_source_properties should be provided"
)
return model
[docs]
class DataSourcesSpecificationConfig(BaseCradleComponentConfig):
"""
Configuration for data sources specification with three-tier field classification.
Corresponds to com.amazon.secureaisandboxproxyservice.models.datasourcesspecification.DataSourcesSpecification
Fields are organized into three tiers:
1. Tier 1: Essential User Inputs - fields that users must explicitly provide
2. Tier 2: System Inputs with Defaults - fields with reasonable defaults that can be overridden
3. Tier 3: Derived Fields - fields calculated from other fields (private attributes with properties)
"""
# ===== Essential User Inputs (Tier 1) =====
# These are fields that users must explicitly provide
start_date: str = Field(
description="Start timestamp exactly 'YYYY-mm-DDTHH:MM:SS', e.g. '2025-01-01T00:00:00'"
)
end_date: str = Field(
description="End timestamp exactly 'YYYY-mm-DDTHH:MM:SS', e.g. '2025-04-17T00:00:00'"
)
data_sources: List[DataSourceConfig] = Field(
description="List of DataSourceConfig objects (both MDS and EDX)"
)
# ===== System Inputs with Defaults (Tier 2) =====
# None currently for this class
# ===== Derived Fields (Tier 3) =====
# None currently for this class
[docs]
class JobSplitOptionsConfig(BaseCradleComponentConfig):
"""
Corresponds to com.amazon.secureaisandboxproxyservice.models.jobsplitoptions.JobSplitOptions:
- split_job: bool
- days_per_split: int
- merge_sql: str
Fields are organized into three tiers:
1. Tier 1: Essential User Inputs - fields that users must explicitly provide
2. Tier 2: System Inputs with Defaults - fields with reasonable defaults that can be overridden
3. Tier 3: Derived Fields - fields calculated from other fields (private attributes with properties)
"""
# ===== Essential User Inputs (Tier 1) =====
# These are fields that users must explicitly provide (when split_job=True)
merge_sql: Optional[str] = Field(
default=None,
description="SQL to run after merging split results (if split_job=True). "
"For example: 'SELECT * FROM INPUT'.",
)
# ===== System Inputs with Defaults (Tier 2) =====
# These are fields with reasonable defaults that users can override
split_job: bool = Field(
default=False,
description="Whether to split the Cradle job into multiple daily runs",
)
days_per_split: int = Field(
default=7, description="Number of days per split (only used if split_job=True)"
)
[docs]
@field_validator("days_per_split")
@classmethod
def days_must_be_positive(cls, v: int) -> int:
if v < 1:
raise ValueError("days_per_split must be ≥ 1")
return v
[docs]
@model_validator(mode="after")
@classmethod
def require_merge_sql_if_split(
cls, model: "JobSplitOptionsConfig"
) -> "JobSplitOptionsConfig":
if model.split_job and not model.merge_sql:
raise ValueError("If split_job=True, merge_sql must be provided")
return model
# ===== Derived Fields (Tier 3) =====
# None currently for this class
[docs]
class OutputSpecificationConfig(BaseCradleComponentConfig):
"""
Corresponds to com.amazon.secureaisandboxproxyservice.models.outputspecification.OutputSpecification:
- output_schema: List[str]
- output_format: str (e.g. 'PARQUET', 'CSV', etc.)
- output_save_mode: str (e.g. 'ERRORIFEXISTS', 'OVERWRITE', 'APPEND', 'IGNORE')
- output_file_count: int (0 means "auto")
- keep_dot_in_output_schema: bool
- include_header_in_s3_output: bool
Fields are organized into three tiers:
1. Tier 1: Essential User Inputs - fields that users must explicitly provide
2. Tier 2: System Inputs with Defaults - fields with reasonable defaults that can be overridden
3. Tier 3: Derived Fields - fields calculated from other fields (private attributes with properties)
"""
# ===== Essential User Inputs (Tier 1) =====
# These are fields that users must explicitly provide
output_schema: List[str] = Field(
description="List of column names to emit (e.g. ['objectId','transactionDate',…])."
)
# ===== System Inputs with Defaults (Tier 2) =====
# These are fields with reasonable defaults that users can override
job_type: str = Field(
default="training",
description="One of ['training','validation','testing','calibration'] to indicate which dataset this job is pulling",
)
# Pipeline S3 location - needed for output_path calculation
pipeline_s3_loc: Optional[str] = Field(
default=None,
description="S3 location for pipeline artifacts (inherited from parent config)",
)
output_format: str = Field(
default="PARQUET",
description="Format for Cradle output: one of ['CSV','UNESCAPED_TSV','JSON','ION','PARQUET']",
)
output_save_mode: str = Field(
default="ERRORIFEXISTS",
description="One of ['ERRORIFEXISTS','OVERWRITE','APPEND','IGNORE']",
)
output_file_count: int = Field(
default=0, ge=0, description="Number of output files (0 means auto‐split)"
)
keep_dot_in_output_schema: bool = Field(
default=False,
description="If False, replace '.' with '__DOT__' in the output header",
)
include_header_in_s3_output: bool = Field(
default=True, description="Whether to write the header row in S3 output files"
)
# ===== Derived Fields (Tier 3) =====
# These are fields calculated from other fields
_output_path: Optional[str] = PrivateAttr(default=None)
@property
def output_path(self) -> str:
"""Get output path derived from pipeline_s3_loc and job_type."""
if self._output_path is None:
# Use the explicitly provided pipeline_s3_loc field if available
if self.pipeline_s3_loc:
self._output_path = f"{self.pipeline_s3_loc}/data-load/{self.job_type}"
else:
# Fallback for backward compatibility
self._output_path = f"s3://default-bucket/data-load/{self.job_type}"
return self._output_path
# Property validator to ensure the output_path is a valid S3 URI
[docs]
def validate_output_path(self) -> None:
"""Validate that output_path is a valid S3 URI."""
# Make sure we have pipeline_s3_loc set before validation
if not hasattr(self, "pipeline_s3_loc") or not self.pipeline_s3_loc:
# Don't try to validate without pipeline_s3_loc - it will use default
return
if not self.output_path.startswith("s3://"):
raise ValueError("output_path must start with 's3://'")
[docs]
@field_validator("job_type")
@classmethod
def validate_job_type(cls, v: str) -> str:
if not v.replace("_", "").isalnum() or v != v.lower():
raise ValueError(
f"job_type must be lowercase alphanumeric (with underscores), got '{v}'"
)
return v
[docs]
@field_validator("output_save_mode")
@classmethod
def validate_save_mode(cls, v: str) -> str:
# Passed verbatim to the SAIS Cradle SDK, which expects the canonical
# UPPERCASE enum. Accept case-insensitive input but normalize to the
# canonical uppercase value.
allowed = {"ERRORIFEXISTS", "OVERWRITE", "APPEND", "IGNORE"}
v_upper = v.upper()
if v_upper not in allowed:
raise ValueError(
f"output_save_mode must be one of {sorted(allowed)} (case-insensitive), got '{v}'"
)
return v_upper
[docs]
class CradleJobSpecificationConfig(BaseCradleComponentConfig):
"""
Corresponds to com.amazon.secureaisandboxproxyservice.models.cradlejobspecification.CradleJobSpecification:
- cluster_type: str (e.g. 'SMALL', 'MEDIUM', 'LARGE')
- cradle_account: str
- extra_spark_job_arguments: Optional[str]
- job_retry_count: int
Fields are organized into three tiers:
1. Tier 1: Essential User Inputs - fields that users must explicitly provide
2. Tier 2: System Inputs with Defaults - fields with reasonable defaults that can be overridden
3. Tier 3: Derived Fields - fields calculated from other fields (private attributes with properties)
"""
# ===== Essential User Inputs (Tier 1) =====
# These are fields that users must explicitly provide
cradle_account: str = Field(
description="Cradle account name (e.g. 'Buyer-Abuse-RnD-Dev')"
)
# ===== System Inputs with Defaults (Tier 2) =====
# These are fields with reasonable defaults that users can override
cluster_type: str = Field(
default="STANDARD",
description="Cluster size for Cradle job (e.g. 'STANDARD', 'SMALL', 'MEDIUM', 'LARGE')",
)
extra_spark_job_arguments: Optional[str] = Field(
default="", description="Any extra Spark driver options (string or blank)"
)
job_retry_count: int = Field(
default=1, ge=0, description="Number of times to retry on failure (default=1)"
)
[docs]
@field_validator("cluster_type")
@classmethod
def validate_cluster_type(cls, v: str) -> str:
# cluster_type is passed verbatim to the SAIS Cradle SDK
# (CradleJobSpecification), which expects the canonical UPPERCASE enum.
# Accept case-insensitive input but normalize to the canonical uppercase
# value so the SDK always receives a valid enum.
allowed = {"STANDARD", "SMALL", "MEDIUM", "LARGE"}
v_upper = v.upper()
if v_upper not in allowed:
raise ValueError(
f"cluster_type must be one of {sorted(allowed)} (case-insensitive), got '{v}'"
)
return v_upper
[docs]
class CradleDataLoadingConfig(BasePipelineConfig):
"""
Top‐level configuration for creating a CreateCradleDataLoadJobRequest with three-tier field classification.
Fields are organized into three tiers:
1. Tier 1: Essential User Inputs - fields that users must explicitly provide
2. Tier 2: System Inputs with Defaults - fields with reasonable defaults that can be overridden
3. Tier 3: Derived Fields - fields calculated from other fields (private attributes with properties)
This class inherits from BasePipelineConfig (not BaseCradleComponentConfig) to maintain
consistency with other pipeline configurations.
"""
# ===== Essential User Inputs (Tier 1) =====
# These are fields that users must explicitly provide
job_type: str = Field(
description="One of ['training','validation','testing','calibration'] to indicate which dataset this job is pulling"
)
data_sources_spec: DataSourcesSpecificationConfig = Field(
description="Full data‐sources specification (start/end dates plus list of sources)."
)
transform_spec: TransformSpecificationConfig = Field(
description="Transform specification: SQL + job‐split options."
)
output_spec: OutputSpecificationConfig = Field(
description="Output specification: schema, output format, save mode, etc."
)
cradle_job_spec: CradleJobSpecificationConfig = Field(
description="Cradle job specification: cluster type, account, retry count, etc."
)
# ===== System Inputs with Defaults (Tier 2) =====
# These are fields with reasonable defaults that users can override
s3_input_override: Optional[str] = Field(
default=None,
description="If set, skip Cradle data pull and use this S3 prefix directly",
)
# ===== Derived Fields (Tier 3) =====
# These are initialized in the model_validator based on other fields
# The output_path in output_spec is a derived field that depends on job_type
# Model configuration - inherit from BasePipelineConfig.Config
model_config = ConfigDict(
arbitrary_types_allowed=True,
validate_assignment=True,
extra="allow", # Allow extra fields like __model_type__ and __model_module__ for type-aware serialization
)
[docs]
@field_validator("job_type")
@classmethod
def validate_job_type(cls, v: str) -> str:
if not v.replace("_", "").isalnum() or v != v.lower():
raise ValueError(
f"job_type must be lowercase alphanumeric (with underscores), got '{v}'"
)
return v
[docs]
@model_validator(mode="after")
def initialize_derived_fields(self) -> "CradleDataLoadingConfig":
"""Initialize all derived fields once after validation."""
# Initialize base class derived fields first
super().initialize_derived_fields()
# Override the output_spec job_type with the parent's job_type
# This ensures consistency between parent and child configs
self.output_spec.job_type = self.job_type
# Pass the pipeline_s3_loc to output_spec for output_path calculation
if hasattr(self, "pipeline_s3_loc"):
self.output_spec.pipeline_s3_loc = self.pipeline_s3_loc
return self
[docs]
def categorize_fields(self) -> Dict[str, List[str]]:
"""
Categorize all fields into three tiers:
1. Tier 1: Essential User Inputs - fields with no defaults (required)
2. Tier 2: System Inputs - fields with defaults (optional)
3. Tier 3: Derived Fields - properties that access private attributes
Returns:
Dict with keys 'essential', 'system', and 'derived' mapping to lists of field names
"""
# Initialize categories
categories = {
"essential": [], # Tier 1: Required, public
"system": [], # Tier 2: Optional (has default), public
"derived": [], # Tier 3: Public properties
}
# Get model fields
model_fields = self.__class__.model_fields
# Categorize public fields into essential (required) or system (with defaults)
for field_name, field_info in model_fields.items():
# Skip private fields
if field_name.startswith("_"):
continue
# Use is_required() to determine if a field is essential
if field_info.is_required():
categories["essential"].append(field_name)
else:
categories["system"].append(field_name)
# Find derived properties (public properties that aren't in model_fields)
for attr_name in dir(self):
if (
not attr_name.startswith("_")
and attr_name not in model_fields
and isinstance(getattr(type(self), attr_name, None), property)
):
categories["derived"].append(attr_name)
# Add nested derived fields
if hasattr(self.output_spec, "categorize_fields"):
nested_categories = self.output_spec.categorize_fields()
# Add the nested derived fields with a prefix
for nested_field in nested_categories["derived"]:
categories["derived"].append(f"output_spec.{nested_field}")
return categories
[docs]
def get_all_tiered_fields(self) -> Dict[str, List[str]]:
"""
Get a flattened list of all fields (including nested fields)
organized by tier.
Returns:
Dict with keys 'essential', 'system', and 'derived' mapping to
lists of field names with dot notation for nesting
"""
return get_flattened_fields(self)
[docs]
def check_split_and_override(self) -> None:
"""Check consistency of split settings and overrides."""
# If splitting is enabled, merge_sql must be provided
if (
self.transform_spec.job_split_options.split_job
and not self.transform_spec.job_split_options.merge_sql
):
raise ValueError("When split_job=True, merge_sql must be provided")