"""DynamoDB clip_meta_heavy asset for processing DynamoDB exports to Snowflake.""" from datetime import datetime import warnings from pathlib import Path import dagster as dg from dagster_snowflake import SnowflakeResource from src.utils.automation_conditions import hourly_cron_condition from src.utils.snowflake.constants import TIME_WINDOW_FRESHNESS_POLICY_WARN_1H_FAIL_2H, PartitionExpr, Warehouse, SnowflakeDB, Team, SnowflakeSchema, Group from src.utils.snowflake.logger import log_query from src.utils.snowflake.query import load_query from src.utils.glue_utils import trigger_glue_job warnings.filterwarnings("ignore", category=dg.BetaWarning) # Get directory of this file for relative SQL file loading ASSET_DIR = Path(__file__).parent DDB_CLIP_META_HEAVY_START_DATE = datetime.strptime('2025-01-01', '%Y-%m-%d') DDB_CLIP_META_HEAVY_TABLE_NAME = "DDB_CLIP_META_HEAVY" WAREHOUSE = Warehouse.DYNAMODB_EVENTS_SMALL.value SNOWFLAKE_DB = SnowflakeDB.SUNO_PROD.value SNOWFLAKE_SCHEMA = SnowflakeSchema.PROD.value @dg.asset( name="ddb_clip_meta_heavy", description="DynamoDB clip_meta_heavy table with heavy text fields from S3 exports.", group_name=Group.DYNAMODB.value, partitions_def=dg.HourlyPartitionsDefinition(start_date=DDB_CLIP_META_HEAVY_START_DATE, end_offset=0), backfill_policy=dg.BackfillPolicy.multi_run(max_partitions_per_run=24*7), automation_condition=hourly_cron_condition, owners=[Team.CORE_POD.value], metadata={ "database": SNOWFLAKE_DB, "schema": SNOWFLAKE_SCHEMA, "table_name": DDB_CLIP_META_HEAVY_TABLE_NAME, "data_start_date": DDB_CLIP_META_HEAVY_START_DATE.strftime("%Y-%m-%d"), "cluster_by": "[p_date, p_hour, id]", "partition_expr": PartitionExpr.HOURLY.value, "transient": True, "sla_minutes": 60, }, freshness_policy=TIME_WINDOW_FRESHNESS_POLICY_WARN_1H_FAIL_2H, ) def ddb_clip_meta_heavy(context: dg.AssetExecutionContext, snowflake: SnowflakeResource) -> dg.MaterializeResult: """ Process DynamoDB clip_meta_heavy data from S3 exports and load into Snowflake. Steps: 1. Trigger Glue job to export DynamoDB data to S3 2. Create table if not exists 3. Delete existing data for the partition window 4. Upsert data from S3 stage """ run_id = context.run.run_id logger = dg.get_dagster_logger() # Get partition time window for processing partition_start = context.partition_time_window.start partition_end = context.partition_time_window.end # Build S3 path for this partition # Format: @SUNO_DYNAMODB_EVENTS/clip-meta-heavy/pdate=YYYY-MM-DD/phour=HH/AWSDynamoDB/data stage_path = f"@SUNO_DYNAMODB_EVENTS/clip-meta-heavy/pdate={partition_start.strftime('%Y-%m-%d')}/phour={partition_start.strftime('%H')}/AWSDynamoDB/data" fetch_params = { "partition_start_date": partition_start.strftime("%Y-%m-%d"), "partition_end_date": partition_end.strftime("%Y-%m-%d"), "partition_start_hour": partition_start.hour, "partition_end_hour": partition_end.hour, "ddb_clip_meta_heavy_table_name": DDB_CLIP_META_HEAVY_TABLE_NAME, "stage_path": stage_path, "p_date": partition_start.strftime("%Y-%m-%d"), "p_hour": partition_start.hour, } logger.info(f"Processing ddb_clip_meta_heavy for partition: {partition_start} to {partition_end}") logger.info(f"Stage path: {stage_path}") logger.info(f"Fetch params: {fetch_params}") rows_affected = 0 glue_job_status = "N/A" glue_job_id = "N/A" glue_execution_time = 0 # Step 1: Trigger Glue job to export DynamoDB data to S3 logger.info("Step 1: Triggering Glue job to export DynamoDB clip_meta_heavy data to S3...") try: glue_result = trigger_glue_job( job_name="ddb_to_s3_clip_meta_heavy_hourly_insert", context=context, arguments={ "--partition_date": partition_start.strftime("%Y-%m-%d"), "--partition_hour": str(partition_start.hour), }, poll_interval=20, wait_for_completion=True, timeout=1800, # 30 minutes ) logger.info(f"Glue job completed successfully: {glue_result}") glue_job_status = "SUCCESS" glue_job_id = glue_result.get("job_run_id", "N/A") glue_execution_time = glue_result.get("execution_time", 0) except Exception as e: logger.error(f"Glue job failed: {str(e)}") glue_job_status = "FAILED" glue_job_id = "N/A" glue_execution_time = 0 # Re-raise the exception to fail the asset materialization raise e # Step 2: Process data in Snowflake with snowflake.get_connection() as conn: cursor = conn.cursor() logger.info(f"Using warehouse {WAREHOUSE}") warehouse_query = load_query( "src/utils/snowflake/queries/use_warehouse.sql", params={"warehouse": WAREHOUSE} ) log_query(logger, warehouse_query) cursor.execute(warehouse_query) # Step 2: Delete existing data for the partition window logger.info(f"Step 2: Deleting existing data from {DDB_CLIP_META_HEAVY_TABLE_NAME}...") delete_query = load_query( "src/utils/snowflake/queries/delete_hourly_partitions.sql", params={**fetch_params, "delete_partition_table_name": DDB_CLIP_META_HEAVY_TABLE_NAME} ) log_query(logger, delete_query) cursor.execute(delete_query) # Step 3: Upsert data from S3 stage using upsert.sql logger.info(f"Step 3: Upserting data into {DDB_CLIP_META_HEAVY_TABLE_NAME} from S3 stage.") upsert_query = load_query( ASSET_DIR / "upsert.sql", params={ "ddb_clip_meta_heavy_table_name": DDB_CLIP_META_HEAVY_TABLE_NAME, "stage_path": stage_path, "p_date": fetch_params["p_date"], "p_hour": fetch_params["p_hour"], } ) log_query(logger, upsert_query) cursor.execute(upsert_query) rows_affected = cursor.rowcount logger.info(f"Successfully processed partition. Affected {rows_affected} rows.") return dg.MaterializeResult( metadata={ "run_id": dg.MetadataValue.text(run_id), "table_name": DDB_CLIP_META_HEAVY_TABLE_NAME, "partition_time_window_start": dg.MetadataValue.text(partition_start.isoformat()), "partition_time_window_end": dg.MetadataValue.text(partition_end.isoformat()), "rows_affected": rows_affected, "stage_path": dg.MetadataValue.text(stage_path), "glue_job_status": dg.MetadataValue.text(glue_job_status), "glue_job_run_id": dg.MetadataValue.text(glue_job_id), "glue_execution_time_seconds": glue_execution_time, }, )