From fa363a33c7af7e88cd216cd6f4dfc64691b13df7 Mon Sep 17 00:00:00 2001
From: Ethan Dalool <git@voussoir.net>
Date: Wed, 31 Mar 2021 17:24:57 -0700
Subject: [PATCH] Improve handling of livestreams, premieres with
 live_broadcast attr.

---
 frontends/ycdl_flask/templates/channel.html |  4 +++
 utilities/database_upgrader.py              | 37 ++++++++++++++++++++
 ycdl/constants.py                           |  3 +-
 ycdl/objects.py                             | 38 ++++++++++++++-------
 ycdl/ycdldb.py                              | 21 +++++++++++-
 ycdl/ytapi.py                               |  3 ++
 6 files changed, 92 insertions(+), 14 deletions(-)

diff --git a/frontends/ycdl_flask/templates/channel.html b/frontends/ycdl_flask/templates/channel.html
index dc1346b..460ab3e 100644
--- a/frontends/ycdl_flask/templates/channel.html
+++ b/frontends/ycdl_flask/templates/channel.html
@@ -226,6 +226,7 @@ https://stackoverflow.com/a/35153397
                 onclick="return action_button_passthrough(event, api.videos.mark_state, 'pending');"
                 >Revert to Pending</button>
 
+                {% if video.live_broadcast is none %}
                 <button
                 {% if video.state == "pending" %}
                 class="video_action_download"
@@ -234,6 +235,9 @@ https://stackoverflow.com/a/35153397
                 {% endif %}
                 onclick="return action_button_passthrough(event, api.videos.start_download);"
                 >Download</button>
+                {% else %}
+                <button disabled>Video is {{video.live_broadcast}}</button>
+                {% endif %}
 
                 <button
                 {% if video.state == "pending" %}
diff --git a/utilities/database_upgrader.py b/utilities/database_upgrader.py
index 27c3420..c8e4608 100644
--- a/utilities/database_upgrader.py
+++ b/utilities/database_upgrader.py
@@ -249,6 +249,43 @@ def upgrade_7_to_8(ycdldb):
     # Redundant due to (state, published)
     ycdldb.sql.execute('DROP INDEX IF EXISTS index_video_state')
 
+def upgrade_8_to_9(ycdldb):
+    '''
+    In this version, the `live_broadcast` column was added to the videos table.
+    '''
+    m = Migrator(ycdldb)
+
+    m.tables['videos']['create'] = '''
+    CREATE TABLE IF NOT EXISTS videos(
+        id TEXT,
+        published INT,
+        author_id TEXT,
+        title TEXT,
+        description TEXT,
+        duration INT,
+        views INT,
+        thumbnail TEXT,
+        live_broadcast TEXT,
+        state TEXT
+    );
+    '''
+    m.tables['videos']['transfer'] = '''
+    INSERT INTO videos SELECT
+        id,
+        published,
+        author_id,
+        title,
+        description,
+        duration,
+        views,
+        thumbnail,
+        NULL,
+        state
+    FROM videos_old;
+    '''
+
+    m.go()
+
 def upgrade_all(data_directory):
     '''
     Given the directory containing a ycdl database, apply all of the
diff --git a/ycdl/constants.py b/ycdl/constants.py
index 37dc42a..9a7f117 100644
--- a/ycdl/constants.py
+++ b/ycdl/constants.py
@@ -1,6 +1,6 @@
 from voussoirkit import sqlhelpers
 
-DATABASE_VERSION = 8
+DATABASE_VERSION = 9
 DB_VERSION_PRAGMA = f'''
 PRAGMA user_version = {DATABASE_VERSION};
 '''
@@ -32,6 +32,7 @@ CREATE TABLE IF NOT EXISTS videos(
     duration INT,
     views INT,
     thumbnail TEXT,
+    live_broadcast TEXT,
     state TEXT
 );
 
diff --git a/ycdl/objects.py b/ycdl/objects.py
index 465a0ea..712bdab 100644
--- a/ycdl/objects.py
+++ b/ycdl/objects.py
@@ -99,20 +99,33 @@ class Channel(Base):
             if not (force or status['new']):
                 break
 
+        # Now we will refresh some other IDs that may not have been refreshed
+        # by the previous loop.
+        refresh_ids = set()
+
+        # 1. Videos which have become unlisted, therefore not returned by the
+        # get_playlist_videos call. Take the set of all known ids minus those
+        # refreshed by the earlier  loop, the difference will be unlisted,
+        # private, or deleted videos. At this time we have no special handling
+        # for deleted videos, but they simply won't come back from ytapi.
         if force:
-            # If some videos have become unlisted, then they will not have been
-            # refreshed by the previous loop. So, take the set of all known ids
-            # minus those refreshed by the loop, and try to refresh them.
-            # Of course, it's possible they were deleted.
             known_ids = {v.id for v in self.ycdldb.get_videos(channel_id=self.id)}
-            refresh_ids = list(known_ids.difference(seen_ids))
-            if refresh_ids:
-                self.ycdldb.log.debug(
-                    '%d ids did not come back from the generator, fetching them separately.',
-                    len(refresh_ids),
-                )
-            for video in self.ycdldb.youtube.get_videos(refresh_ids):
-                self.ycdldb.insert_video(video, commit=False)
+            refresh_ids.update(known_ids.difference(seen_ids))
+
+        # 2. Premieres or live events which may now be over but were not
+        # included in the requested batch of IDs because they are not the most
+        # recent.
+        query = 'SELECT * FROM videos WHERE live_broadcast IS NOT NULL'
+        videos = self.ycdldb.get_videos_by_sql(query)
+        refresh_ids.update(v.id for v in videos)
+
+        if refresh_ids:
+            self.ycdldb.log.debug('Refreshing %d ids separately.', len(refresh_ids))
+
+        # We call ingest_video instead of insert_video so that
+        # premieres / livestreams which have finished can be automarked.
+        for video_id in self.ycdldb.youtube.get_videos(refresh_ids):
+            self.ycdldb.ingest_video(video_id, commit=False)
 
         if commit:
             self.ycdldb.commit()
@@ -186,6 +199,7 @@ class Video(Base):
         self.duration = db_row['duration']
         self.views = db_row['views']
         self.thumbnail = db_row['thumbnail']
+        self.live_broadcast = db_row['live_broadcast']
         self.state = db_row['state']
 
     def __repr__(self):
diff --git a/ycdl/ycdldb.py b/ycdl/ycdldb.py
index 1aa4a96..304416d 100644
--- a/ycdl/ycdldb.py
+++ b/ycdl/ycdldb.py
@@ -433,6 +433,13 @@ class YCDLDBVideoMixin:
             return status
 
         if author.automark == 'downloaded':
+            if video.live_broadcast is not None:
+                self.log.debug(
+                    'Not downloading %s because live_broadcast=%s.',
+                    video.id,
+                    video.live_broadcast,
+                )
+                return status
             # download_video contains a call to mark_state.
             self.download_video(video.id, commit=False)
         else:
@@ -452,9 +459,11 @@ class YCDLDBVideoMixin:
 
         try:
             existing = self.get_video(video.id)
+            existing_live_broadcast = existing.live_broadcast
             download_status = existing.state
         except exceptions.NoSuchVideo:
             existing = None
+            existing_live_broadcast = None
             download_status = 'pending'
 
         data = {
@@ -466,6 +475,7 @@ class YCDLDBVideoMixin:
             'duration': video.duration,
             'views': video.views,
             'thumbnail': video.thumbnail['url'],
+            'live_broadcast': video.live_broadcast,
             'state': download_status,
         }
 
@@ -484,7 +494,16 @@ class YCDLDBVideoMixin:
         if commit:
             self.commit()
 
-        return {'new': not existing, 'video': video}
+        # For the benefit of ingest_video, which will only apply the channel's
+        # automark to newly released videos, let's consider the video to be
+        # new if live_broadcast has changed to be None since last time.
+        # This way, premieres and livestreams can be automarked by the next
+        # refresh after they've ended.
+        is_new = (
+            (existing is None) or
+            (existing_live_broadcast is not None and video.live_broadcast is None)
+        )
+        return {'new': is_new, 'video': video}
 
 class YCDLDB(
         YCDLDBCacheManagerMixin,
diff --git a/ycdl/ytapi.py b/ycdl/ytapi.py
index fe3f5a4..ffcf597 100644
--- a/ycdl/ytapi.py
+++ b/ycdl/ytapi.py
@@ -32,6 +32,9 @@ class Video:
         # Something like '2016-10-01T21:00:01'
         self.published_string = snippet['publishedAt']
         self.published = isodate.parse_datetime(self.published_string).timestamp()
+        self.live_broadcast = snippet['liveBroadcastContent']
+        if self.live_broadcast == 'none':
+            self.live_broadcast = None
         self.tags = snippet.get('tags', [])
 
         self.duration = isodate.parse_duration(content_details['duration']).seconds