from datetime import datetime, timedelta import dagster as dg import pandas as pd from dagster_snowflake import SnowflakeResource from dagster import EnvVar from src.assets.snowflake.dim.assets import ( user_session_state, agg_user_platform_daily, ) from src.utils.snowflake.constants import MIN_TIMESTAMP from src.utils.snowflake.query import PythonStringSQLFormatter ### NOTE: There is a current Dagster bug where the built-in dagster logger does not show up in the UI when using sensors ### https://github.com/dagster-io/dagster/issues/22278 ### For now, we are using the context.log.info() method to log messages to the UI. ### Once this is fixed, we can switch back to using `logger = dg.get_dagster_logger()`. @dg.sensor( name="dim_user_to_snowflake_tasks_sensor", description="Triggers downstream Snowflake tasks after dim_user materializes in Dagster", minimum_interval_seconds=60, metadata={ "migration_sensor": True } ) def dim_user_downstream_tasks_sensor(context: dg.SensorEvaluationContext, snowflake: SnowflakeResource) -> dg.SensorResult: """ Monitors dim_user materializations and triggers downstream Snowflake tasks to maintain orchestration during migration period. """ logger = context.log # Get database and schema from environment variables database = EnvVar("SNOWFLAKE_DB").get_value() schema = EnvVar("SNOWFLAKE_SCHEMA").get_value() # Get the latest materialization event for dim_user latest_materialization = context.instance.get_latest_materialization_event( asset_key=dg.AssetKey("dim_user") ) if not latest_materialization: logger.info("No materialization events found for dim_user") return dg.SensorResult(skip_reason="No materialization events found for dim_user") # Check if we've already processed this materialization last_processed_run_id = context.cursor current_run_id = latest_materialization.run_id logger.info(f"Sensor cursor (last processed): {last_processed_run_id}") logger.info(f"Current materialization run_id: {current_run_id}") # If cursor is None or empty, or if run_ids don't match, process this materialization if last_processed_run_id and last_processed_run_id == current_run_id: # Already processed this materialization logger.info(f"Already processed materialization for run_id: {current_run_id}") return dg.SensorResult(skip_reason=f"Already processed materialization for run_id: {current_run_id}") # New materialization detected (either cursor is None/empty or run_ids don't match) logger.info(f"Processing new materialization - cursor: {last_processed_run_id}, current: {current_run_id}") # New materialization detected - trigger downstream Snowflake tasks # Get partition key from the materialization event partition_key = "unknown" if hasattr(latest_materialization, 'asset_materialization') and latest_materialization.asset_materialization: partition_key = latest_materialization.asset_materialization.partition or "unknown" elif hasattr(latest_materialization, 'partition_key'): partition_key = latest_materialization.partition_key or "unknown" logger.info(f"New dim_user materialization detected (run_id: {current_run_id}, partition: {partition_key})") logger.info(f"Triggering downstream Snowflake tasks in {database}.{schema}...") tasks_triggered = [] tasks_failed = [] with snowflake.get_connection() as conn: cursor = conn.cursor() # Execute downstream tasks that previously depended on dim_user_insert downstream_tasks = [ "FACT_PLAY_HOURLY_UPSERT" ] for task_name in downstream_tasks: try: full_task_name = f"{database}.{schema}.{task_name}" logger.info(f"Executing task: {full_task_name}") cursor.execute(f"EXECUTE TASK {full_task_name}") tasks_triggered.append(task_name) except Exception as e: logger.error(f"Failed to execute task {task_name}: {str(e)}") tasks_failed.append(task_name) # Update cursor to mark this materialization as processed context.update_cursor(current_run_id) if tasks_failed: logger.warning(f"Successfully triggered {len(tasks_triggered)} tasks: {tasks_triggered}") logger.warning(f"Failed to trigger {len(tasks_failed)} tasks: {tasks_failed}") return dg.SensorResult( skip_reason=f"Successfully triggered {len(tasks_triggered)} tasks, but {len(tasks_failed)} failed", metadata={"tasks_triggered": tasks_triggered, "tasks_failed": tasks_failed} ) else: logger.info(f"✅ Successfully triggered all {len(tasks_triggered)} downstream tasks: {tasks_triggered}") return dg.SensorResult( skip_reason=f"Successfully triggered all {len(tasks_triggered)} downstream tasks", metadata={"tasks_triggered": tasks_triggered} ) @dg.sensor( name="snowflake_new_data_sensor__user_session_state", description="Sensor that sends a materialization event when user_session_state has new data. NOTE: This does not display RunRequests in the UI.", minimum_interval_seconds=300, metadata={ "expected_data_cadence": "hourly" } ) def user_session_state_new_data_sensor(context: dg.SensorEvaluationContext, snowflake: SnowflakeResource) -> dg.SensorResult: """Sensor that monitors user_session_state for new data from the previous hour and triggers the asset.""" logger = context.log python_formatter = PythonStringSQLFormatter() logger.info(f"Running sensor `snowflake_new_data_sensor__user_session_state` for asset {user_session_state.key.to_user_string()}") current_hour = datetime.now().replace(minute=0, second=0, microsecond=0) one_hour_ago = current_hour - timedelta(hours=1) one_hour_ago_partition_key = one_hour_ago.strftime("%Y-%m-%d-%H:00") if context.is_first_tick_since_sensor_start: logger.info("This is the first tick since the sensor was started. Setting last processed partition key to MIN_TIMESTAMP.") last_processed_partition_key = MIN_TIMESTAMP.strftime("%Y-%m-%d-%H:00") elif context.cursor: last_processed_partition_key = context.cursor last_processed_hour = datetime.strptime(last_processed_partition_key, "%Y-%m-%d-%H:00") logger.info(f"Last processed partition: {last_processed_partition_key}") if last_processed_partition_key == one_hour_ago_partition_key: logger.info(f"Skipping sensor evaluation because requested partition from previous hour ({one_hour_ago_partition_key}) is equal to the last processed partition ({last_processed_partition_key}).") return dg.SensorResult(skip_reason=f"Requested partition from previous hour ({one_hour_ago_partition_key}) is equal to the last processed partition ({last_processed_partition_key}).") elif last_processed_hour < one_hour_ago: logger.warning(f"Last processed partition ({last_processed_partition_key}) is from more than one hour ago ({one_hour_ago_partition_key}); a backfill may be needed. Proceeding with most recent hour only.") else: logger.warning(f"Cursor is missing, but sensor has been run before. Running the sensor on the most recent hour's ({one_hour_ago_partition_key}) data only, but a backfill may be needed.") last_processed_partition_key = MIN_TIMESTAMP.strftime("%Y-%m-%d-%H:00") with snowflake.get_connection() as conn: cursor = conn.cursor() # Check if there is any data in the previous hour's partition partition_data_check_query = python_formatter.load( "src/utils/snowflake/queries/check_partition_data_hourly.sql", params={ "table_name": user_session_state.key.to_user_string(), "partition_date": one_hour_ago.strftime("%Y-%m-%d"), "partition_hour": one_hour_ago.hour, }, logger=logger, ) cursor.execute(partition_data_check_query) result = cursor.fetchone() conn.commit() logger.info(f"Result: {result}") new_rows = result[0] if result else 0 if new_rows > 0: logger.info(f"Found {new_rows} new rows from previous hour ({one_hour_ago_partition_key}) in {user_session_state.key.to_user_string()}; emitting AssetMaterialization event.") # Return materialization results to directly mark the asset as materialized return dg.SensorResult( cursor=one_hour_ago_partition_key, skip_reason=f"New rows found, emitting AssetMaterialization event for partition {one_hour_ago_partition_key}.", asset_events=[ dg.AssetMaterialization( asset_key=user_session_state.key, partition=one_hour_ago_partition_key, metadata={ "source": "snowflake_new_data_sensor__user_session_state", "dagster/row_count": new_rows, }, ) ], ) else: logger.info(f"No new rows found from previous hour ({one_hour_ago_partition_key}) in {user_session_state.key.to_user_string()}") return dg.SensorResult(skip_reason=f"No new rows found from previous hour ({one_hour_ago_partition_key}) in {user_session_state.key.to_user_string()}") @dg.sensor( name="snowflake_new_data_sensor__agg_user_platform_daily", description="Sensor that sends a materialization event when agg_user_platform_daily has new data for yesterday. Runs at 6:30 AM UTC daily.", minimum_interval_seconds=3600, # Check every hour metadata={ "expected_data_cadence": "daily", "expected_update_time": "~6:00 AM UTC" } ) def agg_user_platform_daily_new_data_sensor(context: dg.SensorEvaluationContext, snowflake: SnowflakeResource) -> dg.SensorResult: """Sensor that monitors agg_user_platform_daily for new data from yesterday and triggers downstream assets.""" logger = context.log python_formatter = PythonStringSQLFormatter() logger.info(f"Running sensor `snowflake_new_data_sensor__agg_user_platform_daily` for asset {agg_user_platform_daily.key.to_user_string()}") # Calculate yesterday's date (D-1) today = datetime.now().replace(hour=0, minute=0, second=0, microsecond=0) yesterday = today - timedelta(days=1) yesterday_partition_key = yesterday.strftime("%Y-%m-%d") if context.is_first_tick_since_sensor_start: logger.info("This is the first tick since the sensor was started. Setting last processed partition key to MIN_TIMESTAMP.") last_processed_partition_key = MIN_TIMESTAMP.strftime("%Y-%m-%d") elif context.cursor: last_processed_partition_key = context.cursor last_processed_date = datetime.strptime(last_processed_partition_key, "%Y-%m-%d") logger.info(f"Last processed partition: {last_processed_partition_key}") if last_processed_partition_key == yesterday_partition_key: logger.info(f"Skipping sensor evaluation because yesterday's partition ({yesterday_partition_key}) has already been processed.") return dg.SensorResult(skip_reason=f"Yesterday's partition ({yesterday_partition_key}) has already been processed.") elif last_processed_date < yesterday: logger.warning(f"Last processed partition ({last_processed_partition_key}) is older than yesterday ({yesterday_partition_key}); a backfill may be needed. Proceeding with yesterday only.") else: logger.warning(f"Cursor is missing, but sensor has been run before. Running the sensor on yesterday's data ({yesterday_partition_key}) only, but a backfill may be needed.") last_processed_partition_key = MIN_TIMESTAMP.strftime("%Y-%m-%d") with snowflake.get_connection() as conn: cursor = conn.cursor() # Check if there is any data in yesterday's partition partition_data_check_query = python_formatter.load( "src/utils/snowflake/queries/check_partition_data_daily.sql", params={ "table_name": f"suno_prod.prod.{agg_user_platform_daily.metadata['table_name']}", "partition_date": yesterday.strftime("%Y-%m-%d"), }, logger=logger, ) cursor.execute(partition_data_check_query) result = cursor.fetchone() conn.commit() logger.info(f"Result: {result}") new_rows = result[0] if result else 0 if new_rows > 0: logger.info(f"Found {new_rows} rows for yesterday ({yesterday_partition_key}) in {agg_user_platform_daily.key.to_user_string()}; emitting AssetMaterialization event.") # Return materialization results to directly mark the asset as materialized return dg.SensorResult( cursor=yesterday_partition_key, skip_reason=f"New rows found, emitting AssetMaterialization event for partition {yesterday_partition_key}.", asset_events=[ dg.AssetMaterialization( asset_key=agg_user_platform_daily.key, partition=yesterday_partition_key, metadata={ "source": "snowflake_new_data_sensor__agg_user_platform_daily", "dagster/row_count": new_rows, }, ) ], ) else: logger.info(f"No data found for yesterday ({yesterday_partition_key}) in {agg_user_platform_daily.key.to_user_string()}") return dg.SensorResult(skip_reason=f"No data found for yesterday ({yesterday_partition_key}) in {agg_user_platform_daily.key.to_user_string()}")