e621_export_downloader 0.0.1

data/lib/e621_export_downloader/export.rb

@@ -0,0 +1,128 @@
+# frozen_string_literal: true
+# typed: strict
+
+require("csv")
+
+module E621ExportDownloader
+  class Export
+    extend(T::Sig)
+    extend(T::Generic)
+
+    Model = type_member
+
+    sig { returns(Date) }
+    attr_accessor(:date)
+
+    sig { returns(ExportHelper[Model]) }
+    attr_accessor(:helper)
+
+    sig { returns(T.nilable(T::Boolean)) }
+    # If nil, no check has been performed yet
+    attr_accessor(:downloaded)
+
+    sig { params(date: T.any(Date, DateTime), helper: ExportHelper[Model]).void }
+    def initialize(date:, helper:)
+      @date = T.let(date.is_a?(DateTime) ? date.to_date : date, Date)
+      @helper = helper
+      @downloaded = T.let(nil, T.nilable(T::Boolean))
+    end
+
+    sig { returns(T::Boolean) }
+    def delete
+      return false unless check_downloaded
+      FileUtils.rm(file_path)
+      @downloaded = false
+      true
+    end
+
+    sig { returns(String) }
+    def download
+      raise(ResolveError, "Export #{helper.type} for #{helper.format_date(date)} does not exist") unless exists?
+      if check_downloaded
+        helper.client.debug("using cached export for #{helper.type.serialize}", header: %W[export #{helper.format_date(date)}])
+        return file_path
+      end
+
+      helper.client.debug("downloading export for #{helper.type.serialize}", header: %W[export #{helper.format_date(date)}])
+
+      FileUtils.mkdir_p(Constants::TEMP_DIR)
+      File.open(file_path, "wb") do |file|
+        inflater = Zlib::Inflate.new(Zlib::MAX_WBITS + 16)
+
+        res = helper.client.connection.get("#{file_name}.gz") do |req|
+          req.options.on_data = proc do |chunk, _total|
+            decompressed = inflater.inflate(chunk)
+            file.write(decompressed) if decompressed && !decompressed.empty?
+          end
+        end
+
+        file.write(inflater.finish)
+        inflater.close
+        file.close
+
+        unless res.success?
+          raise(ResolveError, "Failed to download export #{helper.type.serialize} for #{helper.format_date(date)}: #{res.status} #{res.reason_phrase}")
+        end
+      rescue # rubocop:disable Style/RescueStandardError
+        FileUtils.rm_f(file_path)
+        raise
+      end
+
+      @downloaded = true
+      file_path
+    end
+
+    sig { returns(T::Boolean) }
+    def exists?
+      helper.client.connection.head("#{file_name}.gz").success?
+    end
+
+    sig { params(block: T.proc.params(arg0: Model, arg1: Integer).void).void }
+    def read(&block)
+      download unless check_downloaded
+      helper.client.debug("reading export for #{helper.type.serialize}", header: %W[export #{helper.format_date(date)}])
+      total = line_count
+      CSV.foreach(file_path, headers: true) do |row|
+        args = [helper.parser.call(T.cast(row, CSV::Row).to_hash), total]
+        args = args.slice(0, block.arity) if block.arity != -1
+        block.call(*T.unsafe(args))
+      end
+    end
+
+    sig { returns(T::Array[Model]) }
+    def read_all
+      download unless check_downloaded
+      helper.client.debug("reading all records for #{helper.type.serialize}", header: %W[export #{helper.format_date(date)}])
+      results = []
+      read do |record|
+        results << record
+      end
+      results
+    end
+
+    private
+
+    sig { returns(Integer) }
+    def line_count
+      total = 0
+      CSV.foreach(file_path) { total += 1 }
+      total
+    end
+
+    sig { returns(T::Boolean) }
+    def check_downloaded
+      return @downloaded unless @downloaded.nil?
+      @downloaded = File.exist?(file_path)
+    end
+
+    sig { returns(String) }
+    def file_path
+      File.join(Constants::TEMP_DIR, file_name)
+    end
+
+    sig { returns(String) }
+    def file_name
+      "#{helper.type.serialize}-#{helper.format_date(date)}.csv"
+    end
+  end
+end

data/lib/e621_export_downloader/export_helper.rb

@@ -0,0 +1,83 @@
+# frozen_string_literal: true
+# typed: strict
+
+module E621ExportDownloader
+  class ExportHelper
+    extend(T::Sig)
+    extend(T::Generic)
+
+    Model = type_member
+
+    sig { returns(Types) }
+    attr_reader(:type)
+
+    sig { returns(Client::Options::Parser) }
+    attr_reader(:parser)
+
+    sig { returns(Client) }
+    attr_reader(:client)
+
+    sig { params(type: Types, parser: Client::Options::Parser, client: Client).void }
+    def initialize(type:, parser:, client:)
+      @type = type
+      @parser = parser
+      @client = client
+      @rewind_count = T.let(0, Integer)
+    end
+
+    sig { params(date: T.any(Date, DateTime)).returns(String) }
+    def format_date(date)
+      date.strftime("%Y-%m-%d")
+    end
+
+    sig { params(date: T.any(Date, DateTime)).returns(T::Boolean) }
+    def delete(date)
+      client.debug("deleting export for #{type.serialize}", header: %W[helper #{format_date(date)}])
+      _get(date).delete
+    end
+
+    sig { params(date: T.any(Date, DateTime)).returns(String) }
+    def download(date)
+      client.debug("downloading export for #{type.serialize}", header: %W[helper #{format_date(date)}])
+      _get(date).download
+    end
+
+    sig { params(date: T.any(Date, DateTime)).returns(T::Boolean) }
+    def exists?(date)
+      client.debug("checking export existence for #{type.serialize}", header: %W[helper #{format_date(date)}])
+      _get(date).exists?
+    end
+
+    sig { params(date: T.any(Date, DateTime)).returns(Export[Model]) }
+    def get(date)
+      client.debug("creating export handle for #{type.serialize}", header: %W[helper #{format_date(date)}])
+      Export.new(date: date, helper: self)
+    end
+
+    private
+
+    sig { returns(Integer) }
+    def max_rewind_count
+      return 0 if client.options.rewind_on_not_found == false
+      return 2 if client.options.rewind_on_not_found == true
+      T.cast(client.options.rewind_on_not_found, Integer)
+    end
+
+    sig { params(date: T.any(Date, DateTime), original_date: T.any(Date, DateTime)).returns(Export[Model]) }
+    def _get(date, original_date = date)
+      export = get(date)
+      if export.exists?
+        client.debug("resolved export for #{type.serialize}", header: %w[helper])
+        return export
+      end
+
+      if @rewind_count < max_rewind_count
+        @rewind_count += 1
+        client.debug("rewinding export lookup for #{type.serialize} from #{format_date(original_date)} (attempt #{@rewind_count}/#{max_rewind_count})", header: %w[helper])
+        return _get(date.to_date + 1, original_date)
+      end
+
+      raise(ResolveError, "Export #{type.serialize} for #{format_date(original_date)} does not exist, and either rewinding is not allowed or the maximum rewind limit has been reached")
+    end
+  end
+end

data/lib/e621_export_downloader/models/pool.rb

@@ -0,0 +1,69 @@
+# frozen_string_literal: true
+# typed: strong
+
+module E621ExportDownloader
+  module Models
+    class Pool
+      extend(T::Sig)
+
+      sig { returns(T::Hash[String, String]) }
+      attr_reader(:record)
+
+      sig { returns(String) }
+      attr_reader(:category)
+
+      sig { returns(DateTime) }
+      attr_reader(:created_at)
+
+      sig { returns(Integer) }
+      attr_reader(:creator_id)
+
+      sig { returns(String) }
+      attr_reader(:description)
+
+      sig { returns(Integer) }
+      attr_reader(:id)
+
+      sig { returns(T::Boolean) }
+      attr_reader(:is_active)
+
+      sig { returns(String) }
+      attr_reader(:name)
+
+      sig { returns(T::Array[Integer]) }
+      attr_reader(:post_ids)
+
+      sig { returns(T.nilable(DateTime)) }
+      attr_reader(:updated_at)
+
+      sig { params(record: T::Hash[String, String]).void }
+      def initialize(record)
+        @record = T.let(record, T::Hash[String, String])
+        @category = T.let(T.must(record["category"]), String)
+        @created_at = T.let(DateTime.parse(record["created_at"]), DateTime)
+        @creator_id = T.let(record["creator_id"].to_i, Integer)
+        @description = T.let(T.must(record["description"]), String)
+        @id = T.let(record["id"].to_i, Integer)
+        @is_active = T.let(record["is_active"] == "t", T::Boolean)
+        @name = T.let(T.must(record["name"]), String)
+        @post_ids = T.let(T.must(T.must(record["post_ids"])[1..-2]).split(",").map(&:to_i), T::Array[Integer])
+        @updated_at = T.let(T.must(record["updated_at"]).empty? ? nil : DateTime.parse(record["updated_at"]), T.nilable(DateTime))
+      end
+
+      sig { params(_args: T.untyped).returns(String) }
+      def to_json(*_args)
+        {
+          category:    @category,
+          created_at:  @created_at,
+          creator_id:  @creator_id,
+          description: @description,
+          id:          @id,
+          is_active:   @is_active,
+          name:        @name,
+          post_ids:    @post_ids,
+          updated_at:  @updated_at,
+        }.to_json
+      end
+    end
+  end
+end

data/lib/e621_export_downloader/models/post.rb

@@ -0,0 +1,166 @@
+# frozen_string_literal: true
+# typed: strict
+
+module E621ExportDownloader
+  module Models
+    class Post
+      extend(T::Sig)
+
+      sig { returns(T.nilable(Integer)) }
+      attr_reader(:approver_id)
+
+      sig { returns(Integer) }
+      attr_reader(:change_seq)
+
+      sig { returns(Integer) }
+      attr_reader(:comment_count)
+
+      sig { returns(DateTime) }
+      attr_reader(:created_at)
+
+      sig { returns(String) }
+      attr_reader(:description)
+
+      sig { returns(Integer) }
+      attr_reader(:down_score)
+
+      sig { returns(T.nilable(Float)) }
+      attr_reader(:duration)
+
+      sig { returns(Integer) }
+      attr_reader(:fav_count)
+
+      sig { returns(String) }
+      attr_reader(:file_ext)
+
+      sig { returns(Integer) }
+      attr_reader(:file_size)
+
+      sig { returns(Integer) }
+      attr_reader(:id)
+
+      sig { returns(Integer) }
+      attr_reader(:image_height)
+
+      sig { returns(Integer) }
+      attr_reader(:image_width)
+
+      sig { returns(T::Boolean) }
+      attr_reader(:is_deleted)
+
+      sig { returns(T::Boolean) }
+      attr_reader(:is_flagged)
+
+      sig { returns(T::Boolean) }
+      attr_reader(:is_note_locked)
+
+      sig { returns(T::Boolean) }
+      attr_reader(:is_pending)
+
+      sig { returns(T::Boolean) }
+      attr_reader(:is_rating_locked)
+
+      sig { returns(T::Boolean) }
+      attr_reader(:is_status_locked)
+
+      sig { returns(String) }
+      attr_reader(:locked_tags)
+
+      sig { returns(T.nilable(String)) }
+      attr_reader(:md5)
+
+      sig { returns(T.nilable(Integer)) }
+      attr_reader(:parent_id)
+
+      sig { returns(String) }
+      attr_reader(:rating)
+
+      sig { returns(Integer) }
+      attr_reader(:score)
+
+      sig { returns(T::Array[String]) }
+      attr_reader(:sources)
+
+      sig { returns(T::Array[String]) }
+      attr_reader(:tags)
+
+      sig { returns(Integer) }
+      attr_reader(:up_score)
+
+      sig { returns(T.nilable(DateTime)) }
+      attr_reader(:updated_at)
+
+      sig { returns(T.nilable(Integer)) }
+      attr_reader(:uploader_id)
+
+      sig { params(record: T::Hash[String, String]).void }
+      def initialize(record)
+        @record = T.let(record, T::Hash[String, String])
+        @approver_id = T.let(T.must(record["approver_id"]).empty? ? nil : record["approver_id"].to_i, T.nilable(Integer))
+        @change_seq = T.let(record["change_seq"].to_i, Integer)
+        @comment_count = T.let(record["comment_count"].to_i, Integer)
+        @created_at = T.let(DateTime.parse(record["created_at"]), DateTime)
+        @description = T.let(T.must(record["description"]).gsub("\r\n", "\n"), String)
+        @down_score = T.let(record["down_score"].to_i, Integer)
+        @duration = T.let(T.must(record["duration"]).empty? ? nil : record["duration"].to_f, T.nilable(Float))
+        @fav_count = T.let(record["fav_count"].to_i, Integer)
+        @file_ext = T.let(T.must(record["file_ext"]), String)
+        @file_size = T.let(record["file_size"].to_i, Integer)
+        @id = T.let(record["id"].to_i, Integer)
+        @image_height = T.let(record["image_height"].to_i, Integer)
+        @image_width = T.let(record["image_width"].to_i, Integer)
+        @is_deleted = T.let(record["is_deleted"] == "t", T::Boolean)
+        @is_flagged = T.let(record["is_flagged"] == "t", T::Boolean)
+        @is_note_locked = T.let(record["is_note_locked"] == "t", T::Boolean)
+        @is_pending = T.let(record["is_pending"] == "t", T::Boolean)
+        @is_rating_locked = T.let(record["is_rating_locked"] == "t", T::Boolean)
+        @is_status_locked = T.let(record["is_rating_locked"] == "t", T::Boolean)
+        @locked_tags = T.let(T.must(record["locked_tags"]), String)
+        @md5 = T.let(T.must(record["md5"]).empty? ? nil : record["md5"], T.nilable(String))
+        @parent_id = T.let(T.must(record["parent_id"]).empty? ? nil : record["parent_id"].to_i, T.nilable(Integer))
+        @rating = T.let(T.must(record["rating"]), String)
+        @score = T.let(record["score"].to_i, Integer)
+        @sources = T.let(T.must(record["source"]).gsub("\r\n", "\n").split("\n"), T::Array[String])
+        @tags = T.let(T.must(record["tag_string"]).split, T::Array[String])
+        @up_score = T.let(record["up_score"].to_i, Integer)
+        @updated_at = T.let(T.must(record["updated_at"]).empty? ? nil : DateTime.parse(record["updated_at"]), T.nilable(DateTime))
+        @uploader_id = T.let(T.must(record["uploader_id"]).empty? ? nil : record["uploader_id"].to_i, T.nilable(Integer))
+      end
+
+      sig { params(_args: T.untyped).returns(String) }
+      def to_json(*_args)
+        {
+          approver_id:      @approver_id,
+          change_seq:       @change_seq,
+          comment_count:    @comment_count,
+          created_at:       @created_at,
+          description:      @description,
+          down_score:       @down_score,
+          duration:         @duration,
+          fav_count:        @fav_count,
+          file_ext:         @file_ext,
+          file_size:        @file_size,
+          id:               @id,
+          image_height:     @image_height,
+          image_width:      @image_width,
+          is_deleted:       @is_deleted,
+          is_flagged:       @is_flagged,
+          is_note_locked:   @is_note_locked,
+          is_pending:       @is_pending,
+          is_rating_locked: @is_rating_locked,
+          is_status_locked: @is_status_locked,
+          locked_tags:      @locked_tags,
+          md5:              @md5,
+          parent_id:        @parent_id,
+          rating:           @rating,
+          score:            @score,
+          sources:          @sources,
+          tags:             @tags,
+          up_score:         @up_score,
+          updated_at:       @updated_at,
+          uploader_id:      @uploader_id,
+        }.to_json
+      end
+    end
+  end
+end

data/lib/e621_export_downloader/models/tag.rb

@@ -0,0 +1,41 @@
+# frozen_string_literal: true
+# typed: strict
+
+module E621ExportDownloader
+  module Models
+    class Tag
+      extend(T::Sig)
+
+      sig { returns(String) }
+      attr_reader(:category)
+
+      sig { returns(Integer) }
+      attr_reader(:id)
+
+      sig { returns(String) }
+      attr_reader(:name)
+
+      sig { returns(Integer) }
+      attr_reader(:post_count)
+
+      sig { params(record: T::Hash[String, String]).void }
+      def initialize(record)
+        @record = T.let(record, T::Hash[String, String])
+        @category = T.let(T.must(record["category"]), String)
+        @id = T.let(record["id"].to_i, Integer)
+        @name = T.let(T.must(record["name"]), String)
+        @post_count = T.let(record["post_count"].to_i, Integer)
+      end
+
+      sig { params(_args: T.untyped).returns(String) }
+      def to_json(*_args)
+        {
+          category:   @category,
+          id:         @id,
+          name:       @name,
+          post_count: @post_count,
+        }.to_json
+      end
+    end
+  end
+end

data/lib/e621_export_downloader/models/tag_alias.rb

@@ -0,0 +1,46 @@
+# frozen_string_literal: true
+# typed: strict
+
+module E621ExportDownloader
+  module Models
+    class TagAlias
+      extend(T::Sig)
+
+      sig { returns(String) }
+      attr_reader(:antecedent_name)
+
+      sig { returns(String) }
+      attr_reader(:consequent_name)
+
+      sig { returns(T.nilable(DateTime)) }
+      attr_reader(:created_at)
+
+      sig { returns(Integer) }
+      attr_reader(:id)
+
+      sig { returns(String) }
+      attr_reader(:status)
+
+      sig { params(record: T::Hash[String, String]).void }
+      def initialize(record)
+        @record = T.let(record, T::Hash[String, String])
+        @antecedent_name = T.let(T.must(record["antecedent_name"]), String)
+        @consequent_name = T.let(T.must(record["consequent_name"]), String)
+        @created_at = T.let(T.must(record["created_at"]).empty? ? nil : DateTime.parse(record["created_at"]), T.nilable(DateTime))
+        @id = T.let(record["id"].to_i, Integer)
+        @status = T.let(T.must(record["status"]), String)
+      end
+
+      sig { params(_args: T.untyped).returns(String) }
+      def to_json(*_args)
+        {
+          antecedent_name: @antecedent_name,
+          consequent_name: @consequent_name,
+          created_at:      @created_at,
+          id:              @id,
+          status:          @status,
+        }.to_json
+      end
+    end
+  end
+end

data/lib/e621_export_downloader/models/tag_implication.rb

@@ -0,0 +1,46 @@
+# frozen_string_literal: true
+# typed: strict
+
+module E621ExportDownloader
+  module Models
+    class TagImplication
+      extend(T::Sig)
+
+      sig { returns(String) }
+      attr_reader(:antecedent_name)
+
+      sig { returns(String) }
+      attr_reader(:consequent_name)
+
+      sig { returns(T.nilable(DateTime)) }
+      attr_reader(:created_at)
+
+      sig { returns(Integer) }
+      attr_reader(:id)
+
+      sig { returns(String) }
+      attr_reader(:status)
+
+      sig { params(record: T::Hash[String, String]).void }
+      def initialize(record)
+        @record = T.let(record, T::Hash[String, String])
+        @antecedent_name = T.let(T.must(record["antecedent_name"]), String)
+        @consequent_name = T.let(T.must(record["consequent_name"]), String)
+        @created_at = T.let(T.must(record["created_at"]).empty? ? nil : DateTime.parse(record["created_at"]), T.nilable(DateTime))
+        @id = T.let(record["id"].to_i, Integer)
+        @status = T.let(T.must(record["status"]), String)
+      end
+
+      sig { params(_args: T.untyped).returns(String) }
+      def to_json(*_args)
+        {
+          antecedent_name: @antecedent_name,
+          consequent_name: @consequent_name,
+          created_at:      @created_at,
+          id:              @id,
+          status:          @status,
+        }.to_json
+      end
+    end
+  end
+end

data/lib/e621_export_downloader/models/wiki_page.rb

@@ -0,0 +1,61 @@
+# frozen_string_literal: true
+# typed: strict
+
+module E621ExportDownloader
+  module Models
+    class WikiPage
+      extend(T::Sig)
+
+      sig { returns(String) }
+      attr_reader(:body)
+
+      sig { returns(DateTime) }
+      attr_reader(:created_at)
+
+      sig { returns(T.nilable(Integer)) }
+      attr_reader(:creator_id)
+
+      sig { returns(Integer) }
+      attr_reader(:id)
+
+      sig { returns(T::Boolean) }
+      attr_reader(:is_locked)
+
+      sig { returns(String) }
+      attr_reader(:title)
+
+      sig { returns(T.nilable(DateTime)) }
+      attr_reader(:updated_at)
+
+      sig { returns(T.nilable(Integer)) }
+      attr_reader(:uploader_id)
+
+      sig { params(record: T::Hash[String, String]).void }
+      def initialize(record)
+        @record = T.let(record, T::Hash[String, String])
+        @body = T.let(T.must(record["body"]).gsub("\r\n", "\n"), String)
+        @created_at = T.let(DateTime.parse(record["created_at"]), DateTime)
+        @creator_id = T.let(T.must(record["creator_id"]).empty? ? nil : record["creator_id"].to_i, T.nilable(Integer))
+        @id = T.let(record["id"].to_i, Integer)
+        @is_locked = T.let(record["is_locked"] == "t", T::Boolean)
+        @title = T.let(T.must(record["title"]), String)
+        @updated_at = T.let(T.must(record["updated_at"]).empty? ? nil : DateTime.parse(record["updated_at"]), T.nilable(DateTime))
+        @uploader_id = T.let(T.must(record["uploader_id"]).empty? ? nil : record["uploader_id"].to_i, T.nilable(Integer))
+      end
+
+      sig { params(_args: T.untyped).returns(String) }
+      def to_json(*_args)
+        {
+          body:        @body,
+          created_at:  @created_at,
+          creator_id:  @creator_id,
+          id:          @id,
+          is_locked:   @is_locked,
+          title:       @title,
+          updated_at:  @updated_at,
+          uploader_id: @uploader_id,
+        }.to_json
+      end
+    end
+  end
+end

data/lib/e621_export_downloader/types.rb

@@ -0,0 +1,14 @@
+# frozen_string_literal: true
+
+module E621ExportDownloader
+  class Types < T::Enum
+    enums do
+      Pools = new
+      Posts = new
+      TagAliases = new
+      TagImplications = new
+      Tags = new
+      WikiPages = new
+    end
+  end
+end

data/lib/e621_export_downloader/version.rb

@@ -0,0 +1,10 @@
+# frozen_string_literal: true
+# typed: false
+
+# loaded by bundler
+module E621ExportDownloader
+  module Constants
+    VERSION = "0.0.1"
+    WEBSITE = "https://github.com/DonovanDMC/E621ExportDownloader.rb"
+  end
+end

data/lib/e621_export_downloader.rb

@@ -0,0 +1,12 @@
+# frozen_string_literal: true
+
+require("zeitwerk")
+require("sorbet-runtime")
+
+loader = Zeitwerk::Loader.for_gem
+loader.ignore("#{__dir__}/e621_export_downloader/version.rb")
+loader.setup
+
+module E621ExportDownloader
+  class ResolveError < StandardError; end
+end

data/sorbet/config

@@ -0,0 +1,5 @@
+--dir
+.
+--ignore=tmp/
+--ignore=vendor/
+--parser=prism

data/sorbet/rbi/annotations/.gitattributes

@@ -0,0 +1 @@
+**/*.rbi linguist-vendored=true