From bbc5913b8dacaed4d466bcc466a0782d5e3f5edc Mon Sep 17 00:00:00 2001
From: Fijxu <fijxu@nadeko.net>
Date: Tue, 18 Feb 2025 18:11:05 -0300
Subject: [PATCH] feat(database): support for either Redis or PostgreSQL for
 video cache

---
 src/invidious.cr                 | 19 ++++++---
 src/invidious/database/videos.cr | 72 ++++++++++++++++++++++++++------
 2 files changed, 73 insertions(+), 18 deletions(-)

diff --git a/src/invidious.cr b/src/invidious.cr
index ef7db7b1..06ca02fb 100644
--- a/src/invidious.cr
+++ b/src/invidious.cr
@@ -76,12 +76,8 @@ end
 
 HMAC_KEY = CONFIG.hmac_key
 
-PG_DB    = DB.open CONFIG.database_url
-REDIS_DB = Redis::PooledClient.new(unixsocket: CONFIG.redis_socket || nil, url: CONFIG.redis_url || nil)
+PG_DB = DB.open CONFIG.database_url
 
-if REDIS_DB.ping
-  puts "Connected to redis"
-end
 ARCHIVE_URL     = URI.parse("https://archive.org")
 PUBSUB_URL      = URI.parse("https://pubsubhubbub.appspot.com")
 REDDIT_URL      = URI.parse("https://www.reddit.com")
@@ -156,6 +152,19 @@ end
 OUTPUT = CONFIG.output.upcase == "STDOUT" ? STDOUT : File.open(CONFIG.output, mode: "a")
 LOGGER = Invidious::LogHandler.new(OUTPUT, CONFIG.log_level, CONFIG.colorize_logs)
 
+REDIS_DB = begin
+  LOGGER.info "Connecting to Redis compatible DB"
+  redis = Redis::PooledClient.new(unixsocket: CONFIG.redis_socket || nil, url: CONFIG.redis_url || nil)
+  if redis.ping
+    LOGGER.info "Connected to Redis compatible DB via unix domain socket at '#{CONFIG.redis_socket}'" if CONFIG.redis_socket
+    LOGGER.info "Connected to Redis compatible DB via TCP socket at '#{CONFIG.redis_url}'" if CONFIG.redis_url
+  end
+  redis
+rescue ex
+  LOGGER.error "Failed to connect to a Redis compatible DB. Invidious will store the video cache on the PostgresSQL DB"
+  nil
+end
+
 # Check table integrity
 Invidious::Database.check_integrity(CONFIG)
 
diff --git a/src/invidious/database/videos.cr b/src/invidious/database/videos.cr
index 6b9b4169..64d18068 100644
--- a/src/invidious/database/videos.cr
+++ b/src/invidious/database/videos.cr
@@ -1,16 +1,70 @@
 require "./base.cr"
 
 module Invidious::Database::Videos
+  module DBCache
+    extend self
+
+    def set(video : Video, expire_time)
+      if redis = REDIS_DB
+        redis.set(video.id, video.info.to_json, expire_time)
+        redis.set(video.id + ":time", video.updated, expire_time)
+      else
+        request = <<-SQL
+          INSERT INTO videos
+          VALUES ($1, $2, $3)
+          ON CONFLICT (id) DO NOTHING
+        SQL
+
+        PG_DB.exec(request, video.id, video.info.to_json, video.updated)
+      end
+    end
+
+    def del(id : String)
+      if redis = REDIS_DB
+        redis.del(id)
+        redis.del(id + ":time")
+      else
+        request = <<-SQL
+          DELETE FROM videos *
+          WHERE id = $1
+        SQL
+
+        PG_DB.exec(request, id)
+      end
+    end
+
+    def get(id : String)
+      if redis = REDIS_DB
+        info = redis.get(id)
+        time = redis.get(id + ":time")
+        if info && time
+          return Video.new({
+            id:      id,
+            info:    JSON.parse(info).as_h,
+            updated: Time.parse(time, "%Y-%m-%d %H:%M:%S %z", Time::Location::UTC),
+          })
+        else
+          return nil
+        end
+      else
+        request = <<-SQL
+          SELECT * FROM videos
+          WHERE id = $1
+        SQL
+
+        return PG_DB.query_one?(request, id, as: Video)
+      end
+    end
+  end
+
   extend self
 
   def insert(video : Video)
-    REDIS_DB.set(video.id, video.info.to_json, ex: 14400)
-    REDIS_DB.set(video.id + ":time", video.updated, ex: 14400)
+    DBCache.set(video: video, expire_time: 14400)
   end
 
   def delete(id)
-    REDIS_DB.del(id)
-    REDIS_DB.del(id + ":time")
+    DBCache.del(id)
   end
 
   def delete_expired
@@ -33,14 +87,6 @@ module Invidious::Database::Videos
   end
 
   def select(id : String) : Video?
-    if ((info = REDIS_DB.get(id)) && (time = REDIS_DB.get(id + ":time")))
-      return Video.new({
-        id:      id,
-        info:    JSON.parse(info).as_h,
-        updated: Time.parse(time, "%Y-%m-%d %H:%M:%S %z", Time::Location::UTC),
-      })
-    else
-      return nil
-    end
+    return DBCache.get(id)
   end
 end