Only retry failed assets and simplify sensor function

Author: mbertrand

dg_projects/learning_resources/learning_resources/sensors/youtube_shorts.py

@@ -191,6 +191,90 @@ def _create_dynamic_partition_requests(
     return dynamic_requests
 
 
+def fetch_and_log_video_metadata(youtube_client, youtube_config_provider):
+    """
+    Fetch video metadata for the 24 most recent videos and log the results.
+
+    Returns:
+        current_video_ids: Set of video IDs
+        metadata_by_video_id: Dictionary of metadata by video ID
+    """
+    current_videos_metadata = fetch_youtube_shorts_api_data(
+        youtube_client, youtube_config_provider
+    )
+
+    current_video_ids = {
+        item["contentDetails"]["videoId"] for item in current_videos_metadata
+    }
+
+    metadata_by_video_id = {
+        item["contentDetails"]["videoId"]: item for item in current_videos_metadata
+    }
+
+    return current_video_ids, metadata_by_video_id
+
+
+def determine_videos_to_process(current_video_ids, successfully_processed):
+    """
+    Determine which videos need processing (new or failed).
+
+    Returns:
+        Set of video IDs to process
+    """
+    videos_to_process = current_video_ids - successfully_processed
+
+    if not videos_to_process:
+        # Assign message to variable for clarity
+        skip_reason_message = (
+            "No videos need processing and none were removed from YouTube"
+        )
+        raise SkipReason(skip_reason_message)
+
+    return videos_to_process
+
+
+def create_run_requests(videos_to_process, metadata_by_video_id):
+    """
+    Create run requests for videos that need processing.
+
+    Returns:
+        List of RunRequest objects
+    """
+    run_requests = []
+    for video_id in videos_to_process:
+        metadata_json = (
+            json.dumps(metadata_by_video_id.get(video_id))
+            if video_id in metadata_by_video_id
+            else None
+        )
+
+        run_requests.append(
+            RunRequest(
+                asset_selection=[
+                    AssetKey(["youtube_shorts", "video_content"]),
+                    AssetKey(["youtube_shorts", "video_thumbnail"]),
+                    AssetKey(["youtube_shorts", "video_metadata"]),
+                    AssetKey(["youtube_shorts", "video_webhook"]),
+                ],
+                partition_key=video_id,
+                tags={"video_id": video_id},
+                run_config={
+                    "ops": {
+                        "youtube_shorts__video_thumbnail": {
+                            "config": {"cached_metadata": metadata_json}
+                        },
+                        "youtube_shorts__video_metadata": {
+                            "config": {"cached_metadata": metadata_json}
+                        },
+                    }
+                }
+                if metadata_json
+                else {},
+            )
+        )
+    return run_requests
+
+
 @sensor(
     description=(
         "Sensor to monitor YouTube channels for new video shorts and cleanup "
@@ -226,22 +310,10 @@ def youtube_shorts_sensor(context):
         youtube_config_provider = context.resources.youtube_config_provider
 
         # Fetch video metadata for the 24 most recent videos
-        current_videos_metadata = fetch_youtube_shorts_api_data(
+        current_video_ids, metadata_by_video_id = fetch_and_log_video_metadata(
             youtube_client, youtube_config_provider
         )
 
-        # Extract video IDs from metadata and create a lookup dict
-        current_video_ids = {
-            item["contentDetails"]["videoId"] for item in current_videos_metadata
-        }
-
-        # Create a metadata lookup dictionary for passing to assets
-        metadata_by_video_id = {
-            item["contentDetails"]["videoId"]: item for item in current_videos_metadata
-        }
-
-        context.log.info("Found %d videos in YouTube channels", len(current_video_ids))
-
         # Get existing dynamic partitions
         existing_partitions = set(
             context.instance.get_dynamic_partitions("youtube_video_ids")
@@ -254,58 +326,12 @@ def youtube_shorts_sensor(context):
         )
 
         # Determine videos that need processing (new or failed)
-        videos_to_process = current_video_ids - successfully_processed
-
-        # Log which specific videos need processing
-        if videos_to_process:
-            sample_to_process = list(videos_to_process)[:10]
-            context.log.info("Sample videos needing processing: %s", sample_to_process)
-
-        if not videos_to_process:
-            return SkipReason(
-                "No videos need processing and none were removed from YouTube"
-            )
-
-        context.log.info(
-            "Videos needing processing (new or failed): %d", len(videos_to_process)
+        videos_to_process = determine_videos_to_process(
+            current_video_ids, successfully_processed
         )
 
-        # Create run requests for videos that need processing (new or failed)
-        run_requests = []
-        total_config_size = 0
-        for video_id in videos_to_process:
-            metadata_json = (
-                json.dumps(metadata_by_video_id.get(video_id))
-                if video_id in metadata_by_video_id
-                else None
-            )
-            if metadata_json:
-                total_config_size += len(metadata_json)
-
-            run_requests.append(
-                RunRequest(
-                    asset_selection=[
-                        AssetKey(["youtube_shorts", "video_content"]),
-                        AssetKey(["youtube_shorts", "video_thumbnail"]),
-                        AssetKey(["youtube_shorts", "video_metadata"]),
-                        AssetKey(["youtube_shorts", "video_webhook"]),
-                    ],
-                    partition_key=video_id,
-                    tags={"video_id": video_id},
-                    run_config={
-                        "ops": {
-                            "youtube_shorts__video_thumbnail": {
-                                "config": {"cached_metadata": metadata_json}
-                            },
-                            "youtube_shorts__video_metadata": {
-                                "config": {"cached_metadata": metadata_json}
-                            },
-                        }
-                    }
-                    if metadata_json
-                    else {},
-                )
-            )
+        # Create run requests for videos that need processing
+        run_requests = create_run_requests(videos_to_process, metadata_by_video_id)
 
         # Create dynamic partition requests
         dynamic_requests = _create_dynamic_partition_requests(
@@ -328,6 +354,10 @@ def youtube_shorts_sensor(context):
             cursor=cursor_json,
         )
 
-    except Exception:
-        context.log.exception("Error in YouTube shorts sensor")
-        return SkipReason("Sensor failed due to error - check logs")
+    except SkipReason as skip:
+        return skip
+
+    except Exception as e:
+        msg = f"Error in youtube_shorts_sensor: {e!s}"
+        context.log.exception(msg)
+        raise