import warnings from datetime import datetime import dagster as dg from dagster import EnvVar from dagster_snowflake import SnowflakeResource # Using asset key references to avoid import chain issues from src.utils.automation_conditions import hourly_cron_with_eager_historical_backfill_condition from src.utils.snowflake.constants import TIME_WINDOW_FRESHNESS_POLICY_WARN_1H_FAIL_2H, PartitionExpr, Warehouse from src.utils.snowflake.query import JinjaSQLFormatter, PythonStringSQLFormatter warnings.filterwarnings("ignore", category=dg.BetaWarning) DIM_HOOK_START_DATE = datetime.strptime('2025-05-12', '%Y-%m-%d') DIM_HOOK_TABLE_NAME = "DIM_HOOK" class DimHookConfig(dg.Config): dim_hook_table_name: str = DIM_HOOK_TABLE_NAME warehouse: str = Warehouse.DIM_HOOK_MEDIUM.value @dg.asset( name="dim_hook", description="Dimension table for hook/video metadata and properties, including AI-generated features from Gemini analysis.", group_name="hooks", partitions_def=dg.HourlyPartitionsDefinition(start_date=DIM_HOOK_START_DATE, end_offset=-1), deps=[ dg.AssetDep("rds_video_hook"), dg.AssetDep("bk_nonrealtime_event"), dg.AssetDep("app_event"), dg.AssetDep("ddb_item_info"), dg.AssetDep(["PROD", "dim_clip"]), ], backfill_policy=dg.BackfillPolicy.multi_run(max_partitions_per_run=24*7), owners=["team:core-pod"], metadata={ "database": EnvVar("SNOWFLAKE_DB").get_value(), "schema": EnvVar("SNOWFLAKE_SCHEMA").get_value(), "table_name": DIM_HOOK_TABLE_NAME, "data_start_date": DIM_HOOK_START_DATE.strftime("%Y-%m-%d"), "cluster_by": "[p_date, p_hour]", "partition_expr": PartitionExpr.HOURLY.value, "transient": True, "sla_minutes": 120, }, automation_condition=hourly_cron_with_eager_historical_backfill_condition, freshness_policy=TIME_WINDOW_FRESHNESS_POLICY_WARN_1H_FAIL_2H, ) def dim_hook(context: dg.AssetExecutionContext, snowflake: SnowflakeResource, config: DimHookConfig) -> dg.MaterializeResult: run_id = context.run.run_id logger = dg.get_dagster_logger() jinja_formatter = JinjaSQLFormatter() python_formatter = PythonStringSQLFormatter() # 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, "dim_hook_table_name": config.dim_hook_table_name, } logger.info(f"Processing dim_hook 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 {config.warehouse}") warehouse_query = python_formatter.load("src/utils/snowflake/queries/use_warehouse.sql", params={"warehouse": config.warehouse}, logger=logger) cursor.execute(warehouse_query) # 1. Ensure the target table exists logger.info(f"Creating table {config.dim_hook_table_name}...") create_table_query = jinja_formatter.load("src/assets/snowflake/dim/hook/table.sql", params=fetch_params, logger=logger) cursor.execute(create_table_query) # 2. Delete existing data from the target table for the partition window logger.info(f"Deleting existing data from {config.dim_hook_table_name} for partition window {partition_start} to {partition_end}.") delete_query = python_formatter.load("src/utils/snowflake/queries/delete_hourly_partitions.sql", params={**fetch_params, "delete_partition_table_name": config.dim_hook_table_name}, logger=logger) cursor.execute(delete_query) # 3. Insert new hook dimension data logger.info(f"Inserting new hook dimension data into {config.dim_hook_table_name}.") insert_query = jinja_formatter.load("src/assets/snowflake/dim/hook/insert_dim_hook.sql", params=fetch_params, logger=logger) cursor.execute(insert_query) # 4. Commit the transaction conn.commit() rows_inserted = cursor.rowcount logger.info(f"Successfully processed partition. Inserted {rows_inserted} rows.") return dg.MaterializeResult( metadata={ "run_id": dg.MetadataValue.text(run_id), "table_name": config.dim_hook_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_inserted if not context.has_partition_key_range else 0, }, ) # dim_hook_hourly_insert = dg.define_asset_job( # name="dim_hook_hourly_insert", # description="Hourly job to insert dim_hook asset data", # selection=dg.AssetSelection.assets(dim_hook), # partitions_def=dg.HourlyPartitionsDefinition(start_date=DIM_HOOK_START_DATE, end_offset=-2), # tags={"team": "core-pod", "domain": "dimensions", "type": "asset-based"}, # ) # # Hourly schedule for the job # dim_hook_hourly_schedule = dg.ScheduleDefinition( # name="dim_hook_hourly_schedule", # cron_schedule="0 * * * *", # Every hour at minute 0 # job=dim_hook_hourly_insert, # description="Schedule to run dim_hook_hourly_insert every hour", # tags={"team": "core-pod", "cadence": "hourly", "domain": "dimensions"}, # )