from datetime import datetime, timedelta 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, PartitionExpr, Warehouse, SnowflakeDB, Team, SnowflakeSchema, Group from src.utils.snowflake.logger import log_query from src.utils.snowflake.query import load_query warnings.filterwarnings("ignore", category=dg.BetaWarning) # Get directory of this file for relative SQL file loading ASSET_DIR = Path(__file__).parent WEB_HOOK_EVENT_START_DATE = datetime.strptime('2024-06-01', '%Y-%m-%d') WEB_HOOK_EVENT_TABLE_NAME = "WEB_HOOK_EVENT" WAREHOUSE = Warehouse.WEB_EVENT_X_SMALL.value SNOWFLAKE_DB = SnowflakeDB.SUNO_PROD.value SNOWFLAKE_SCHEMA = SnowflakeSchema.PROD.value @dg.asset( name="web_hook_event", description="Web hook events for user interactions with hook features.", group_name=Group.SEGMENT_EVENTS.value, partitions_def=dg.HourlyPartitionsDefinition(start_date=WEB_HOOK_EVENT_START_DATE, end_offset=0), backfill_policy=dg.BackfillPolicy.multi_run(max_partitions_per_run=24*7), owners=[Team.CORE_POD.value], metadata={ "database": SNOWFLAKE_DB, "schema": SNOWFLAKE_SCHEMA, "table_name": WEB_HOOK_EVENT_TABLE_NAME, "data_start_date": WEB_HOOK_EVENT_START_DATE.strftime("%Y-%m-%d"), "cluster_by": "[p_date, p_hour]", "partition_expr": PartitionExpr.HOURLY.value, "transient": True, "sla_minutes": 15, }, freshness_policy=TIME_WINDOW_FRESHNESS_POLICY_WARN_1H_FAIL_2H, ) def web_hook_event(context: dg.AssetExecutionContext, snowflake: SnowflakeResource) -> dg.MaterializeResult: 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 is_multi_partition_range = context.has_partition_key_range 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, "web_hook_event_table_name": WEB_HOOK_EVENT_TABLE_NAME, } logger.info(f"Processing web_hook_event for partition: {partition_start} to {partition_end}") logger.info(f"Fetch params: {fetch_params}") 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) # 1. Delete existing data for the partition window logger.info(f"Deleting existing data from {WEB_HOOK_EVENT_TABLE_NAME} for partition window {partition_start} to {partition_end}.") delete_query = load_query("src/utils/snowflake/queries/delete_hourly_partitions.sql", params={**fetch_params, "delete_partition_table_name": WEB_HOOK_EVENT_TABLE_NAME}) log_query(logger, delete_query) cursor.execute(delete_query) # 2. Insert new data hour by hour using insert.sql total_rows_inserted = 0 current_hour = partition_start while current_hour < partition_end: # Generate file path and geo table name for this hour file_path = f"@SUNO_S3_EVENTS/Hook-Web-Event/0/{current_hour.year}/{current_hour.month}/{current_hour.day}/{current_hour.hour}" geo_table_name = f"GEO_INFO_TEMP_HOOK_{current_hour.year}_{current_hour.month}_{current_hour.day}_{current_hour.hour}" logger.info(f"Processing hour: {current_hour.strftime('%Y-%m-%d %H:00:00')}") logger.info(f"File path: {file_path}") logger.info(f"Geo table: {geo_table_name}") # Step 3a: Insert data from S3 using insert.sql insert_params = { "web_hook_event_table_name": WEB_HOOK_EVENT_TABLE_NAME, "file_path": file_path, "partition_date": current_hour.strftime("%Y-%m-%d"), "partition_hour": current_hour.hour, } insert_query = load_query(ASSET_DIR / "insert.sql", params=insert_params) log_query(logger, insert_query) cursor.execute(insert_query) rows_for_hour = cursor.rowcount total_rows_inserted += rows_for_hour logger.info(f"Inserted {rows_for_hour} rows for hour {current_hour.strftime('%Y-%m-%d %H:00:00')}") # Step 3b: Create temporary geo info table logger.info(f"Creating geo info temp table: {geo_table_name}") create_geo_table_query = f""" CALL CREATE_GEO_INFO_TEMP_TABLE_PROC( '{geo_table_name}', (SELECT ARRAY_AGG(DISTINCT client_ip) FROM {WEB_HOOK_EVENT_TABLE_NAME} WHERE p_date = '{current_hour.strftime('%Y-%m-%d')}' AND p_hour = {current_hour.hour}) ); """ log_query(logger, create_geo_table_query) cursor.execute(create_geo_table_query) # Step 3c: Merge geo info back into the table logger.info(f"Merging geo info for hour {current_hour.strftime('%Y-%m-%d %H:00:00')}") geo_merge_params = { "web_hook_event_table_name": WEB_HOOK_EVENT_TABLE_NAME, "geo_table_name": geo_table_name, "p_date": current_hour.strftime('%Y-%m-%d'), "p_hour": current_hour.hour, } geo_merge_query = load_query(ASSET_DIR / "geo_merge.sql", params=geo_merge_params) log_query(logger, geo_merge_query) cursor.execute(geo_merge_query) # Step 3d: Drop the temporary geo table logger.info(f"Dropping geo info temp table: {geo_table_name}") drop_geo_table_query = f"DROP TABLE IF EXISTS {geo_table_name};" log_query(logger, drop_geo_table_query) cursor.execute(drop_geo_table_query) # Move to next hour current_hour += timedelta(hours=1) logger.info(f"Successfully processed partition. Inserted {total_rows_inserted} rows.") return dg.MaterializeResult( metadata={ "run_id": dg.MetadataValue.text(run_id), "table_name": WEB_HOOK_EVENT_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": total_rows_inserted if not is_multi_partition_range else 0, }, )