From 6704ce3214ed8c9a3acb69e0fc53824324a610fc Mon Sep 17 00:00:00 2001 From: Samantaz Fox Date: Mon, 6 Dec 2021 18:03:37 +0100 Subject: [PATCH] Move DB utility functions to the proper module --- src/invidious.cr | 20 ++--- src/invidious/database/base.cr | 106 ++++++++++++++++++++++++ src/invidious/helpers/helpers.cr | 105 ----------------------- src/invidious/jobs/refresh_feeds_job.cr | 2 +- 4 files changed, 117 insertions(+), 116 deletions(-) diff --git a/src/invidious.cr b/src/invidious.cr index 561fc9cf..ef02f143 100644 --- a/src/invidious.cr +++ b/src/invidious.cr @@ -113,19 +113,19 @@ LOGGER = Invidious::LogHandler.new(OUTPUT, CONFIG.log_level) # Check table integrity if CONFIG.check_tables - check_enum(PG_DB, "privacy", PlaylistPrivacy) + Invidious::Database.check_enum(PG_DB, "privacy", PlaylistPrivacy) - check_table(PG_DB, "channels", InvidiousChannel) - check_table(PG_DB, "channel_videos", ChannelVideo) - check_table(PG_DB, "playlists", InvidiousPlaylist) - check_table(PG_DB, "playlist_videos", PlaylistVideo) - check_table(PG_DB, "nonces", Nonce) - check_table(PG_DB, "session_ids", SessionId) - check_table(PG_DB, "users", User) - check_table(PG_DB, "videos", Video) + Invidious::Database.check_table(PG_DB, "channels", InvidiousChannel) + Invidious::Database.check_table(PG_DB, "channel_videos", ChannelVideo) + Invidious::Database.check_table(PG_DB, "playlists", InvidiousPlaylist) + Invidious::Database.check_table(PG_DB, "playlist_videos", PlaylistVideo) + Invidious::Database.check_table(PG_DB, "nonces", Nonce) + Invidious::Database.check_table(PG_DB, "session_ids", SessionId) + Invidious::Database.check_table(PG_DB, "users", User) + Invidious::Database.check_table(PG_DB, "videos", Video) if CONFIG.cache_annotations - check_table(PG_DB, "annotations", Annotation) + Invidious::Database.check_table(PG_DB, "annotations", Annotation) end end diff --git a/src/invidious/database/base.cr b/src/invidious/database/base.cr index 055a6284..6e49ea1a 100644 --- a/src/invidious/database/base.cr +++ b/src/invidious/database/base.cr @@ -1,4 +1,110 @@ require "pg" module Invidious::Database + extend self + + def check_enum(db, enum_name, struct_type = nil) + return # TODO + + if !db.query_one?("SELECT true FROM pg_type WHERE typname = $1", enum_name, as: Bool) + LOGGER.info("check_enum: CREATE TYPE #{enum_name}") + + db.using_connection do |conn| + conn.as(PG::Connection).exec_all(File.read("config/sql/#{enum_name}.sql")) + end + end + end + + def check_table(db, table_name, struct_type = nil) + # Create table if it doesn't exist + begin + db.exec("SELECT * FROM #{table_name} LIMIT 0") + rescue ex + LOGGER.info("check_table: check_table: CREATE TABLE #{table_name}") + + db.using_connection do |conn| + conn.as(PG::Connection).exec_all(File.read("config/sql/#{table_name}.sql")) + end + end + + return if !struct_type + + struct_array = struct_type.type_array + column_array = get_column_array(db, table_name) + column_types = File.read("config/sql/#{table_name}.sql").match(/CREATE TABLE public\.#{table_name}\n\((?[\d\D]*?)\);/) + .try &.["types"].split(",").map(&.strip).reject &.starts_with?("CONSTRAINT") + + return if !column_types + + struct_array.each_with_index do |name, i| + if name != column_array[i]? + if !column_array[i]? + new_column = column_types.select(&.starts_with?(name))[0] + LOGGER.info("check_table: ALTER TABLE #{table_name} ADD COLUMN #{new_column}") + db.exec("ALTER TABLE #{table_name} ADD COLUMN #{new_column}") + next + end + + # Column doesn't exist + if !column_array.includes? name + new_column = column_types.select(&.starts_with?(name))[0] + db.exec("ALTER TABLE #{table_name} ADD COLUMN #{new_column}") + end + + # Column exists but in the wrong position, rotate + if struct_array.includes? column_array[i] + until name == column_array[i] + new_column = column_types.select(&.starts_with?(column_array[i]))[0]?.try &.gsub("#{column_array[i]}", "#{column_array[i]}_new") + + # There's a column we didn't expect + if !new_column + LOGGER.info("check_table: ALTER TABLE #{table_name} DROP COLUMN #{column_array[i]}") + db.exec("ALTER TABLE #{table_name} DROP COLUMN #{column_array[i]} CASCADE") + + column_array = get_column_array(db, table_name) + next + end + + LOGGER.info("check_table: ALTER TABLE #{table_name} ADD COLUMN #{new_column}") + db.exec("ALTER TABLE #{table_name} ADD COLUMN #{new_column}") + + LOGGER.info("check_table: UPDATE #{table_name} SET #{column_array[i]}_new=#{column_array[i]}") + db.exec("UPDATE #{table_name} SET #{column_array[i]}_new=#{column_array[i]}") + + LOGGER.info("check_table: ALTER TABLE #{table_name} DROP COLUMN #{column_array[i]} CASCADE") + db.exec("ALTER TABLE #{table_name} DROP COLUMN #{column_array[i]} CASCADE") + + LOGGER.info("check_table: ALTER TABLE #{table_name} RENAME COLUMN #{column_array[i]}_new TO #{column_array[i]}") + db.exec("ALTER TABLE #{table_name} RENAME COLUMN #{column_array[i]}_new TO #{column_array[i]}") + + column_array = get_column_array(db, table_name) + end + else + LOGGER.info("check_table: ALTER TABLE #{table_name} DROP COLUMN #{column_array[i]} CASCADE") + db.exec("ALTER TABLE #{table_name} DROP COLUMN #{column_array[i]} CASCADE") + end + end + end + + return if column_array.size <= struct_array.size + + column_array.each do |column| + if !struct_array.includes? column + LOGGER.info("check_table: ALTER TABLE #{table_name} DROP COLUMN #{column} CASCADE") + db.exec("ALTER TABLE #{table_name} DROP COLUMN #{column} CASCADE") + end + end + end + + def get_column_array(db, table_name) + column_array = [] of String + db.query("SELECT * FROM #{table_name} LIMIT 0") do |rs| + rs.column_count.times do |i| + column = rs.as(PG::ResultSet).field(i) + column_array << column.name + end + end + + return column_array + end end diff --git a/src/invidious/helpers/helpers.cr b/src/invidious/helpers/helpers.cr index c5f6c6c5..982bcda1 100644 --- a/src/invidious/helpers/helpers.cr +++ b/src/invidious/helpers/helpers.cr @@ -60,111 +60,6 @@ def html_to_content(description_html : String) return description end -def check_enum(db, enum_name, struct_type = nil) - return # TODO - - if !db.query_one?("SELECT true FROM pg_type WHERE typname = $1", enum_name, as: Bool) - LOGGER.info("check_enum: CREATE TYPE #{enum_name}") - - db.using_connection do |conn| - conn.as(PG::Connection).exec_all(File.read("config/sql/#{enum_name}.sql")) - end - end -end - -def check_table(db, table_name, struct_type = nil) - # Create table if it doesn't exist - begin - db.exec("SELECT * FROM #{table_name} LIMIT 0") - rescue ex - LOGGER.info("check_table: check_table: CREATE TABLE #{table_name}") - - db.using_connection do |conn| - conn.as(PG::Connection).exec_all(File.read("config/sql/#{table_name}.sql")) - end - end - - return if !struct_type - - struct_array = struct_type.type_array - column_array = get_column_array(db, table_name) - column_types = File.read("config/sql/#{table_name}.sql").match(/CREATE TABLE public\.#{table_name}\n\((?[\d\D]*?)\);/) - .try &.["types"].split(",").map(&.strip).reject &.starts_with?("CONSTRAINT") - - return if !column_types - - struct_array.each_with_index do |name, i| - if name != column_array[i]? - if !column_array[i]? - new_column = column_types.select(&.starts_with?(name))[0] - LOGGER.info("check_table: ALTER TABLE #{table_name} ADD COLUMN #{new_column}") - db.exec("ALTER TABLE #{table_name} ADD COLUMN #{new_column}") - next - end - - # Column doesn't exist - if !column_array.includes? name - new_column = column_types.select(&.starts_with?(name))[0] - db.exec("ALTER TABLE #{table_name} ADD COLUMN #{new_column}") - end - - # Column exists but in the wrong position, rotate - if struct_array.includes? column_array[i] - until name == column_array[i] - new_column = column_types.select(&.starts_with?(column_array[i]))[0]?.try &.gsub("#{column_array[i]}", "#{column_array[i]}_new") - - # There's a column we didn't expect - if !new_column - LOGGER.info("check_table: ALTER TABLE #{table_name} DROP COLUMN #{column_array[i]}") - db.exec("ALTER TABLE #{table_name} DROP COLUMN #{column_array[i]} CASCADE") - - column_array = get_column_array(db, table_name) - next - end - - LOGGER.info("check_table: ALTER TABLE #{table_name} ADD COLUMN #{new_column}") - db.exec("ALTER TABLE #{table_name} ADD COLUMN #{new_column}") - - LOGGER.info("check_table: UPDATE #{table_name} SET #{column_array[i]}_new=#{column_array[i]}") - db.exec("UPDATE #{table_name} SET #{column_array[i]}_new=#{column_array[i]}") - - LOGGER.info("check_table: ALTER TABLE #{table_name} DROP COLUMN #{column_array[i]} CASCADE") - db.exec("ALTER TABLE #{table_name} DROP COLUMN #{column_array[i]} CASCADE") - - LOGGER.info("check_table: ALTER TABLE #{table_name} RENAME COLUMN #{column_array[i]}_new TO #{column_array[i]}") - db.exec("ALTER TABLE #{table_name} RENAME COLUMN #{column_array[i]}_new TO #{column_array[i]}") - - column_array = get_column_array(db, table_name) - end - else - LOGGER.info("check_table: ALTER TABLE #{table_name} DROP COLUMN #{column_array[i]} CASCADE") - db.exec("ALTER TABLE #{table_name} DROP COLUMN #{column_array[i]} CASCADE") - end - end - end - - return if column_array.size <= struct_array.size - - column_array.each do |column| - if !struct_array.includes? column - LOGGER.info("check_table: ALTER TABLE #{table_name} DROP COLUMN #{column} CASCADE") - db.exec("ALTER TABLE #{table_name} DROP COLUMN #{column} CASCADE") - end - end -end - -def get_column_array(db, table_name) - column_array = [] of String - db.query("SELECT * FROM #{table_name} LIMIT 0") do |rs| - rs.column_count.times do |i| - column = rs.as(PG::ResultSet).field(i) - column_array << column.name - end - end - - return column_array -end - def cache_annotation(db, id, annotations) if !CONFIG.cache_annotations return diff --git a/src/invidious/jobs/refresh_feeds_job.cr b/src/invidious/jobs/refresh_feeds_job.cr index 926c27fa..4b52c959 100644 --- a/src/invidious/jobs/refresh_feeds_job.cr +++ b/src/invidious/jobs/refresh_feeds_job.cr @@ -25,7 +25,7 @@ class Invidious::Jobs::RefreshFeedsJob < Invidious::Jobs::BaseJob spawn do begin # Drop outdated views - column_array = get_column_array(db, view_name) + column_array = Invidious::Database.get_column_array(db, view_name) ChannelVideo.type_array.each_with_index do |name, i| if name != column_array[i]? LOGGER.info("RefreshFeedsJob: DROP MATERIALIZED VIEW #{view_name}")