"""DynamoDB item_info 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_ITEM_INFO_START_DATE = datetime.strptime('2025-01-01', '%Y-%m-%d') DDB_ITEM_INFO_TABLE_NAME = "DDB_ITEM_INFO" WAREHOUSE = Warehouse.DYNAMODB_EVENTS_LARGE.value SNOWFLAKE_DB = SnowflakeDB.SUNO_PROD.value SNOWFLAKE_SCHEMA = SnowflakeSchema.PROD.value @dg.asset( name="ddb_item_info", description="DynamoDB item_info table from S3 exports.", group_name=Group.DYNAMODB.value, partitions_def=dg.HourlyPartitionsDefinition(start_date=DDB_ITEM_INFO_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_ITEM_INFO_TABLE_NAME, "data_start_date": DDB_ITEM_INFO_START_DATE.strftime("%Y-%m-%d"), "cluster_by": "[p_date, p_hour, type]", "partition_expr": PartitionExpr.HOURLY.value, "transient": True, "sla_minutes": 40, }, freshness_policy=TIME_WINDOW_FRESHNESS_POLICY_WARN_1H_FAIL_2H, ) def ddb_item_info(context: dg.AssetExecutionContext, snowflake: SnowflakeResource) -> dg.MaterializeResult: """ Process DynamoDB item_info 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/item-info/pdate=YYYY-MM-DD/phour=HH/AWSDynamoDB/data stage_path = f"@SUNO_DYNAMODB_EVENTS/item-info/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_item_info_table_name": DDB_ITEM_INFO_TABLE_NAME, "stage_path": stage_path, "p_date": partition_start.strftime("%Y-%m-%d"), "p_hour": partition_start.hour, } logger.info(f"Processing ddb_item_info for partition: {partition_start} to {partition_end}") logger.info(f"Stage path: {stage_path}") logger.info(f"Fetch params: {fetch_params}") rows_affected = 0 # Step 1: Trigger Glue job to export DynamoDB data to S3 logger.info("Step 1: Triggering Glue job to export DynamoDB item_info data to S3...") try: glue_result = trigger_glue_job( job_name="ddb_to_s3_item_info_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=2400, # 40 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_ITEM_INFO_TABLE_NAME}...") delete_query = load_query( "src/utils/snowflake/queries/delete_hourly_partitions.sql", params={**fetch_params, "delete_partition_table_name": DDB_ITEM_INFO_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_ITEM_INFO_TABLE_NAME} from S3 stage.") upsert_query = load_query( ASSET_DIR / "upsert.sql", params={ "ddb_item_info_table_name": DDB_ITEM_INFO_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_ITEM_INFO_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, }, )