# from datetime import timedelta from typing import Any, Mapping, Optional import hashlib import yaml import dagster as dg from dagster import AutomationCondition, DagsterInvalidDefinitionError # from dagster._utils.schedules import is_valid_cron_string # from dagster.preview.freshness import FreshnessPolicy from dagster._utils.tags import is_valid_tag_key from dagster_dbt import DagsterDbtTranslator import pandas as pd from src.utils.automation_conditions import DBT_AUTOMATION_CONDITION_MAP from src.utils.common import str2bool from src.utils.dbt import DBT_PARTITION_MAPPING_MAP, DBT_MIN_PARTITION_DATE class CustomDagsterDbtTranslator(DagsterDbtTranslator): def get_automation_condition(self, dbt_resource_props: Mapping[str, Any]) -> Optional[AutomationCondition]: """ Get automation condition for dbt model. ⚠️ IMPORTANT: Tag-based automation condition configuration is DEPRECATED. Use meta.dagster.automation_condition instead. See DAGSTER_DBT_GUIDE.md for details. If no automation condition is found, it will return the default automation condition. Example model yml file (CORRECT - preferred method): ```yaml models: - name: table_name meta: dagster: automation_condition: daily_cron_with_eager_historical_backfill_condition ``` ❌ DEPRECATED - Do NOT use: ```yaml config: tags: ["automation_condition=daily_cron_with_eager_historical_backfill_condition"] ``` """ # Read automation conditions from config # Preferred method: Read automation conditions directly from config automation_condition = dbt_resource_props.get("meta", {}).get("dagster", {}).get("automation_condition", None) # Legacy method: Read automation conditions from tags tags = dbt_resource_props.get("tags", []) tag_automation_conditions = [tag.split('automation_condition=')[1] for tag in tags if 'automation_condition=' in tag] if automation_condition and tag_automation_conditions: raise DagsterInvalidDefinitionError( f"Both automation condition and tag automation conditions found for model {dbt_resource_props.get('name')}: {automation_condition} and {tag_automation_conditions}. " f"Please use only automation_condition in config; tag automation conditions are deprecated and will be removed in the future." ) # Validate automation_condition if automation_condition: # Only one automation condition is allowed if isinstance(automation_condition, list): raise DagsterInvalidDefinitionError( f"Multiple automation conditions found for model {dbt_resource_props.get('name')}: {automation_condition}. " f"Only one automation condition is allowed. Please combine into one automation condition and pass the combined condition." ) # Validate tag_automation_conditions if tag_automation_conditions: if len(tag_automation_conditions) > 1: raise DagsterInvalidDefinitionError( f"Multiple automation conditions found for model {dbt_resource_props.get('name')}: {tag_automation_conditions}. " f"Only one automation condition is allowed. Please combine into one automation condition and pass the combined condition." ) elif len(tag_automation_conditions) == 1: automation_condition = tag_automation_conditions[0] if automation_condition: if automation_condition in DBT_AUTOMATION_CONDITION_MAP: return DBT_AUTOMATION_CONDITION_MAP[automation_condition] else: raise DagsterInvalidDefinitionError( f"Provided automation condition '{automation_condition}' is not defined in `DBT_AUTOMATION_CONDITION_MAP`. " f"Please add the automation condition to `DBT_AUTOMATION_CONDITION_MAP`.") else: return super().get_automation_condition(dbt_resource_props) def get_tags(self, dbt_resource_props: Mapping[str, Any]) -> Mapping[str, str]: """ Get tags for the dbt model and convert into Dagster tags. dbt tags are strings, whereas Dagster tags are mappings. To support both, tags should be in the format of "key=value". Example model yml file: ```yaml models: - name: table_name tags: ["project=project_name"] ``` """ dbt_tags = dbt_resource_props.get("tags", []) dagster_tags = {} for tag in dbt_tags: if tag.count('=') > 1: raise DagsterInvalidDefinitionError(f"Invalid tag: {tag}. Please use only one '=' in the tag to separate key-value pairs.") key, _, value = tag.partition("=") if value: dagster_tags[key] = value elif is_valid_tag_key(key): dagster_tags[key] = "" else: raise DagsterInvalidDefinitionError(f"Invalid Dagster tag: {tag}") return dagster_tags def get_partitions_def(self, dbt_resource_props: Mapping[str, Any]) -> Optional[dg.PartitionsDefinition]: """ Get the partitions definition for the dbt model and convert into Dagster partitions. Configuration priority: 1. If only tag-based is present → Use tag + meta.dagster.partition_start_date (legacy, deprecated) 2. Otherwise → Use meta.dagster.partitions_def (preferred method) Example model yml file (CORRECT - preferred method): ```yaml models: - name: table_name meta: dagster: partitions_def: type: hourly start_date: 2024-01-01 end_offset: 0 ``` Legacy method (deprecated but supported): ```yaml models: - name: table_name config: tags: ["partition_type=daily"] meta: dagster: partition_start_date: "2024-01-01" end_offset: 0 ``` """ model_name = dbt_resource_props.get('name', 'unknown') dagster_meta = dbt_resource_props.get("meta", {}).get("dagster", {}) tags = dbt_resource_props.get("tags", []) # Check for tag-based partition configuration partition_type_tags = [tag.split('partition_type=')[1] for tag in tags if 'partition_type=' in tag] has_tag_config = len(partition_type_tags) > 0 # Check for meta.dagster.partitions_def partitions_def = dagster_meta.get("partitions_def", None) has_partitions_def = partitions_def is not None # If only tag-based config is present, use legacy method if has_tag_config: partition_type = partition_type_tags[0] legacy_partition_start_date = dagster_meta.get("partition_start_date", None) legacy_end_offset = dagster_meta.get("end_offset", 0) if not legacy_partition_start_date: legacy_partition_start_date = DBT_MIN_PARTITION_DATE start_date = pd.to_datetime(legacy_partition_start_date) if partition_type == "hourly": return dg.HourlyPartitionsDefinition(start_date=start_date, end_offset=legacy_end_offset) elif partition_type == "daily": return dg.DailyPartitionsDefinition(start_date=start_date, end_offset=legacy_end_offset) elif partition_type == "none": return None else: raise DagsterInvalidDefinitionError( f"Model '{model_name}': Unsupported partition type '{partition_type}' from tag. " f"Must be 'daily', 'hourly', or 'none'." ) # Use meta.dagster.partitions_def (preferred method) if has_partitions_def: partition_type = partitions_def.get("type", None) start_date_str = partitions_def.get("start_date", None) end_offset = partitions_def.get("end_offset", 0) if not start_date_str: raise DagsterInvalidDefinitionError( f"Model '{model_name}': start_date is required in meta.dagster.partitions_def" ) if not partition_type: raise DagsterInvalidDefinitionError( f"Model '{model_name}': type is required in meta.dagster.partitions_def. " f"Must be 'daily' or 'hourly'." ) start_date = pd.to_datetime(start_date_str) if partition_type == "hourly": return dg.HourlyPartitionsDefinition(start_date=start_date, end_offset=end_offset) elif partition_type == "daily": return dg.DailyPartitionsDefinition(start_date=start_date, end_offset=end_offset) else: raise DagsterInvalidDefinitionError( f"Model '{model_name}': Unsupported partition type '{partition_type}' in meta.dagster.partitions_def. " f"Must be 'daily' or 'hourly'." ) # No partition configuration found return None def get_partition_mapping(self, dbt_resource_props: Mapping[str, Any], dbt_parent_resource_props: Mapping[str, Any]) -> Optional[dg.PartitionMapping]: """ Get the partition mapping for a given dbt resource and its parent resource. This function is only called on one child-parent resource pair at a time, so it is designed to return one PartitionMapping even if the child resource has multiple parent dependencies. The model yml file should look like this: ```yaml models: - name: table_name meta: dagster: partition_mappings: - asset_key: upstream_table_1_dagster_asset_key type: time_window start_offset: 0 # Optional end_offset: 0 # Optional allow_nonexistent_upstream_partitions: false # Optional ``` """ dbt_child_asset_key = self.get_asset_key(dbt_resource_props) dbt_parent_asset_key = self.get_asset_key(dbt_parent_resource_props) # For self-dependencies, use parent implementation if dbt_child_asset_key == dbt_parent_asset_key: return super().get_partition_mapping(dbt_resource_props, dbt_parent_resource_props) # Check if the current model has partition mapping configuration partition_mappings = dbt_resource_props.get("meta", {}).get("dagster", {}).get("partition_mappings", None) if partition_mappings: # Iterate over dependencies defined in config and look for a matching partition mapping configuration for dagster_dependency in partition_mappings: # Get the dependency asset key from config config_dependency_asset_key = dagster_dependency.get("asset_key", None) if not config_dependency_asset_key: raise DagsterInvalidDefinitionError(f"`asset_key` is required for dependency partition mapping, but was not found for asset {dbt_child_asset_key}") # Convert config asset key to a Dagster AssetKey object # Asset key can be a string (e.g. asset_key: 'dim_clip') or a list (e.g. asset_key: ['prod_int', 'int_daily_active_users']) if isinstance(config_dependency_asset_key, str): dagster_dependency_asset_key = dg.AssetKey([config_dependency_asset_key]) elif isinstance(config_dependency_asset_key, list): dagster_dependency_asset_key = dg.AssetKey(config_dependency_asset_key) # If config asset key does not match the parent asset key in question, move to next dependency listed in config if dagster_dependency_asset_key != dbt_parent_asset_key: continue # Create PartitionMapping object based on the type of partition mapping specified in the config dagster_dependency_partition_mapping_type = dagster_dependency.get("type", None) if dagster_dependency_partition_mapping_type not in DBT_PARTITION_MAPPING_MAP: raise DagsterInvalidDefinitionError( f"Unsupported partition mapping type: {dagster_dependency_partition_mapping_type}. " f"Please use one of the following: [{', '.join(DBT_PARTITION_MAPPING_MAP.keys())}]" ) # TimeWindowPartitionMapping takes parameters; all other supported PartitionMapping types do not if dagster_dependency_partition_mapping_type == "time_window": start_offset = int(dagster_dependency.get("start_offset", 0)) end_offset = int(dagster_dependency.get("end_offset", 0)) allow_nonexistent_upstream_partitions = str2bool(dagster_dependency.get("allow_nonexistent_upstream_partitions", False)) return dg.TimeWindowPartitionMapping( start_offset=start_offset, end_offset=end_offset, allow_nonexistent_upstream_partitions=allow_nonexistent_upstream_partitions ) else: return DBT_PARTITION_MAPPING_MAP[dagster_dependency_partition_mapping_type]() ## TODO: Uncomment this once we have full coverage of partition mappings on dbt models # # Raise warning if no partition mapping is found for given parent # raise DagsterInvalidDefinitionError( # f"ERROR: No partition mapping found for child-parent pair: Child {dbt_child_asset_key} -> Parent {dbt_parent_asset_key}. " # f"Please add a partition mapping to the child dbt node's yml file. " # f"Found dependencies: {partition_mappings}" # ) return super().get_partition_mapping(dbt_resource_props, dbt_parent_resource_props) def get_description(self, dbt_resource_props: Mapping[str, Any]) -> str: """ Get the description for the dbt asset, enhanced with Raw Sql and YAML Configuration sections. This method extends the base description from dbt with: 1. Raw Sql section: Shows the raw SQL code from the dbt model 2. YAML Configuration section: Shows the YAML configuration for the model """ # Get the base description from dbt (from the model's description field in YAML) base_description = dbt_resource_props.get("description", "") # Get raw SQL code raw_code = dbt_resource_props.get("raw_code", "") # Build the enhanced description description_parts = [] # Add base description if it exists if base_description: description_parts.append(base_description) # Add Raw Sql section if raw_code: description_parts.append("\n## Raw Sql\n") description_parts.append("```sql") description_parts.append(raw_code) description_parts.append("```") # Add YAML Configuration section description_parts.append("\n## YAML Configuration\n") description_parts.append("```yaml") # Build YAML dict structure from dbt_resource_props model_name = dbt_resource_props.get("name", "model_name") yaml_dict = { "models": [ { "name": model_name } ] } model_dict = yaml_dict["models"][0] # Add description if exists if base_description: model_dict["description"] = base_description # Add config section config = dbt_resource_props.get("config", {}) if config: model_dict["config"] = config # Dump YAML using yaml.dump() with proper indentation for arrays of objects class CustomDumper(yaml.SafeDumper): def increase_indent(self, flow=False, indentless=False): return super(CustomDumper, self).increase_indent(flow, False) yaml_output = yaml.dump( yaml_dict, Dumper=CustomDumper, default_flow_style=False, sort_keys=False, allow_unicode=True, indent=2, width=1000 # Prevent line wrapping that can break formatting ) description_parts.append(yaml_output) description_parts.append("```") return "\n".join(description_parts) def get_code_version(self, dbt_resource_props: Mapping[str, Any]) -> Optional[str]: """ Override to return a stable code version based on SQL content only. This excludes timestamp metadata that changes on every dbt run, preventing all assets from appearing changed on each deployment. """ # Get the compiled SQL, which is the actual code content compiled_code = dbt_resource_props.get("compiled_code") or dbt_resource_props.get("raw_code", "") # Use the SQL content hash as the code version # This ensures only actual SQL changes trigger a new version if compiled_code: code_hash = hashlib.sha256(compiled_code.encode()).hexdigest() return code_hash # Fallback to parent implementation if no code available return super().get_code_version(dbt_resource_props) # TODO: Uncomment this once Dagster natively supports freshness policies in DagsterDbtTranslator. # https://github.com/dagster-io/dagster/pull/32544 has actually deprecated the get_freshness_policy integration in the DagsterDbtTranslator. # For now, freshness policies are handled inside asset definition factory functions. # def get_freshness_policy(self, dbt_resource_props: Mapping[str, Any]) -> Optional[FreshnessPolicy]: # """ # Get the freshness policy for the dbt model. # This is a custom implementation that uses the preview FreshnessPolicy, not the LegacyFreshnessPolicy class that was deprecated in 1.6. # See: https://docs.dagster.io/guides/observe/asset-freshness-policies # """ # freshness_policy = dbt_resource_props.get("config", {}).get("freshness_policy", None) # if freshness_policy: # freshness_policy_type = freshness_policy.get("type", None) # # Time window freshness policy # # Fields: warn_hours, fail_hours # if freshness_policy_type == "time_window": # warn_hours = freshness_policy.get("warn_hours", None) # fail_hours = freshness_policy.get("fail_hours", None) # if not warn_hours or not fail_hours or not isinstance(warn_hours, int) or not isinstance(fail_hours, int): # raise DagsterInvalidDefinitionError(f"warn_hours (int) and fail_hours (int) are required for time window freshness policy, but got warn_hours: {warn_hours} and fail_hours: {fail_hours}") # return FreshnessPolicy.time_window(warn_window=timedelta(hours=warn_hours), fail_window=timedelta(hours=fail_hours)) # # Cron freshness policy # # Fields: deadline_cron, lower_bound_delta, timezone (optional) # elif freshness_policy_type == "cron": # deadline_cron = freshness_policy.get("deadline_cron", None) # if not deadline_cron or not isinstance(deadline_cron, str): # raise DagsterInvalidDefinitionError(f"deadline_cron (string) is required for cron freshness policy, but got deadline_cron: {deadline_cron}") # if not is_valid_cron_string(deadline_cron): # raise DagsterInvalidDefinitionError(f"Invalid cron string: {deadline_cron}") # lower_bound_delta_hours = freshness_policy.get("lower_bound_delta_hours", None) # if not lower_bound_delta_hours or not isinstance(lower_bound_delta_hours, int): # raise DagsterInvalidDefinitionError(f"lower_bound_delta_hours (int) is required for cron freshness policy, but got lower_bound_delta_hours: {lower_bound_delta_hours}") # lower_bound_delta = timedelta(hours=lower_bound_delta_hours) # timezone = freshness_policy.get("timezone", "UTC") # if not timezone or not isinstance(timezone, str): # raise DagsterInvalidDefinitionError(f"timezone (string) is required for cron freshness policy, but got timezone: {timezone}") # return FreshnessPolicy.cron(deadline_cron=deadline_cron, lower_bound_delta=lower_bound_delta, timezone=timezone) # else: # raise DagsterInvalidDefinitionError(f"Unsupported freshness policy type: {freshness_policy_type}. Please use 'time_window' or 'cron'.") # else: # return None