"""Clip asset for processing RDS clip data and triggering Glue export.""" from datetime import datetime import warnings from pathlib import Path import dagster as dg from dagster_snowflake import SnowflakeResource from src.utils.snowflake.constants import TIME_WINDOW_FRESHNESS_POLICY_WARN_1H_FAIL_2H, Group, PartitionExpr, Warehouse, SnowflakeDB, Team, SnowflakeSchema 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 from src.utils.automation_conditions import long_running_rds_hourly_cron_condition, rds_hourly_cron_condition warnings.filterwarnings("ignore", category=dg.BetaWarning) # Get directory of this file for relative SQL file loading ASSET_DIR = Path(__file__).parent CLIP_START_DATE = datetime.strptime('2024-01-01', '%Y-%m-%d') CLIP_TABLE_NAME = "CLIP" DATABASE_TABLE_NAME = "bots_generatedclip" WAREHOUSE = Warehouse.SUNO_PROD_RDS_CLIP_HOURLY_UPDATE_MEDIUM.value SNOWFLAKE_DB = SnowflakeDB.SUNO_PROD.value SNOWFLAKE_SCHEMA = SnowflakeSchema.PROD.value @dg.asset( name="clip_insert", description="Processed clip table with data from RDS, triggers Glue export to S3.", group_name=Group.RDS.value, partitions_def=dg.HourlyPartitionsDefinition(start_date=CLIP_START_DATE, end_offset=0), backfill_policy=dg.BackfillPolicy.multi_run(max_partitions_per_run=24), owners=[Team.CORE_POD.value], metadata={ "database": SNOWFLAKE_DB, "schema": SNOWFLAKE_SCHEMA, "table_name": CLIP_TABLE_NAME, "data_start_date": CLIP_START_DATE.strftime("%Y-%m-%d"), "cluster_by": "[p_date, p_hour]", "partition_expr": PartitionExpr.HOURLY.value, "transient": True, "sla_minutes": 60, }, automation_condition=long_running_rds_hourly_cron_condition, freshness_policy=TIME_WINDOW_FRESHNESS_POLICY_WARN_1H_FAIL_2H, ) def clip_insert( context: dg.AssetExecutionContext, snowflake: SnowflakeResource ) -> dg.MaterializeResult: """ Process clip data from RDS and trigger Glue export. Steps: 1. Trigger Glue job to export data to S3 2. Create table if not exists 3. Delete existing data for partition 4. Insert 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 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, "clip_table_name": CLIP_TABLE_NAME, "stage_path": f"@SUNO_DATABASE_EVENTS/{DATABASE_TABLE_NAME}_insert/pdate={partition_start.strftime('%Y-%m-%d')}/phour={partition_start.strftime('%H')}", } logger.info(f"Processing clip for partition: {partition_start} to {partition_end}") logger.info(f"Fetch params: {fetch_params}") rows_affected = 0 # Step 1: Trigger Glue job to export data to S3 logger.info("Step 1: Triggering Glue job to export clip data to S3...") try: glue_result = trigger_glue_job( job_name=f"rds_to_s3_{DATABASE_TABLE_NAME}_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=3600, # 60 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) # Delete existing data for the partition window logger.info(f"Deleting existing data from {CLIP_TABLE_NAME}...") delete_query = load_query( "src/utils/snowflake/queries/delete_hourly_partitions.sql", params={**fetch_params, "delete_partition_table_name": CLIP_TABLE_NAME} ) log_query(logger, delete_query) cursor.execute(delete_query) # Insert data from S3 stage logger.info(f"Inserting data into {CLIP_TABLE_NAME} from S3 stage.") insert_query = load_query( ASSET_DIR / "insert.sql", params=fetch_params ) log_query(logger, insert_query) cursor.execute(insert_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": CLIP_TABLE_NAME, "partition_time_window_start": dg.MetadataValue.text(partition_start.isoformat()), "partition_time_window_end": dg.MetadataValue.text(partition_end.isoformat()), "dagster/row_count": rows_affected, "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, }, ) @dg.asset( name="clip_update", description="Upsert clip data from S3 stage, updates existing records and inserts new ones.", group_name=Group.RDS.value, partitions_def=dg.HourlyPartitionsDefinition(start_date=CLIP_START_DATE, end_offset=0), backfill_policy=dg.BackfillPolicy.multi_run(max_partitions_per_run=24), owners=[Team.CORE_POD.value], metadata={ "database": SNOWFLAKE_DB, "schema": SNOWFLAKE_SCHEMA, "table_name": CLIP_TABLE_NAME, "data_start_date": CLIP_START_DATE.strftime("%Y-%m-%d"), "cluster_by": "[p_date, p_hour]", "partition_expr": PartitionExpr.HOURLY.value, "transient": True, "sla_minutes": 60, }, automation_condition=long_running_rds_hourly_cron_condition, freshness_policy=TIME_WINDOW_FRESHNESS_POLICY_WARN_1H_FAIL_2H, ) def clip_update( context: dg.AssetExecutionContext, snowflake: SnowflakeResource ) -> dg.MaterializeResult: """ Upsert clip data from S3 stage into Snowflake. This asset triggers a separate Glue job and performs MERGE operation: - Updates existing records based on ID and partition - Inserts new records that don't exist Steps: 1. Trigger Glue job to export update data to S3 2. Use warehouse 3. Run MERGE query 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 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, "clip_table_name": CLIP_TABLE_NAME, "stage_path": f"@SUNO_DATABASE_EVENTS/{DATABASE_TABLE_NAME}_update/pdate={partition_start.strftime('%Y-%m-%d')}/phour={partition_start.strftime('%H')}", } logger.info(f"Upserting clip data for partition: {partition_start} to {partition_end}") logger.info(f"Fetch params: {fetch_params}") rows_affected = 0 # Step 1: Trigger Glue job to export update data to S3 logger.info("Step 1: Triggering Glue job to export clip update data to S3...") try: glue_result = trigger_glue_job( job_name=f"rds_to_s3_{DATABASE_TABLE_NAME}_hourly_update", 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=3600, # 60 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) # Upsert data from S3 stage logger.info(f"Upserting data into {CLIP_TABLE_NAME} from S3 stage.") upsert_query = load_query( ASSET_DIR / "upsert.sql", params=fetch_params ) log_query(logger, upsert_query) cursor.execute(upsert_query) rows_affected = cursor.rowcount logger.info(f"Successfully upserted partition. Affected {rows_affected} rows.") return dg.MaterializeResult( metadata={ "run_id": dg.MetadataValue.text(run_id), "table_name": CLIP_TABLE_NAME, "partition_time_window_start": dg.MetadataValue.text(partition_start.isoformat()), "partition_time_window_end": dg.MetadataValue.text(partition_end.isoformat()), "dagster/row_count": rows_affected, "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, }, ) @dg.asset( name="clip", description="Final clip asset that completes after both insert and update operations.", group_name=Group.RDS.value, partitions_def=dg.HourlyPartitionsDefinition(start_date=CLIP_START_DATE, end_offset=0), backfill_policy=dg.BackfillPolicy.multi_run(max_partitions_per_run=24), deps=["clip_insert", "clip_update"], owners=[Team.CORE_POD.value], metadata={ "database": SNOWFLAKE_DB, "schema": SNOWFLAKE_SCHEMA, "table_name": CLIP_TABLE_NAME, "data_start_date": CLIP_START_DATE.strftime("%Y-%m-%d"), "cluster_by": "[p_date, p_hour]", "partition_expr": PartitionExpr.HOURLY.value, "transient": True, "sla_minutes": 60, }, automation_condition=dg.AutomationCondition.eager(), freshness_policy=TIME_WINDOW_FRESHNESS_POLICY_WARN_1H_FAIL_2H, ) def clip( context: dg.AssetExecutionContext, snowflake: SnowflakeResource ) -> dg.MaterializeResult: """ Final clip asset that materializes after both insert and update operations complete. This asset: - Depends on clip_insert (initial data load from Glue) - Depends on clip_update (upsert operations) - Validates the final state of the partition - Returns row count statistics """ 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 logger.info(f"Finalizing clip data for partition: {partition_start} to {partition_end}") # Query the final row count for this partition 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) # Get final row count for the partition row_count_query = f""" SELECT COUNT(*) as row_count FROM {SNOWFLAKE_DB}.{SNOWFLAKE_SCHEMA}.{CLIP_TABLE_NAME} WHERE p_date = '{partition_start.strftime('%Y-%m-%d')}' AND p_hour = {partition_start.hour} """ log_query(logger, row_count_query) cursor.execute(row_count_query) result = cursor.fetchone() row_count = result[0] if result else 0 logger.info(f"Final partition has {row_count} rows after insert and update operations.") return dg.MaterializeResult( metadata={ "run_id": dg.MetadataValue.text(run_id), "table_name": CLIP_TABLE_NAME, "partition_time_window_start": dg.MetadataValue.text(partition_start.isoformat()), "partition_time_window_end": dg.MetadataValue.text(partition_end.isoformat()), "dagster/row_count": row_count, }, )