Add compact index support for private sources

.rubocop-bundler.yml

@@ -15,6 +15,9 @@ Lint/AssignmentInCondition:
 Lint/UnusedMethodArgument:
   Enabled: false
 
+Lint/UriEscapeUnescape:
+  Enabled: true
+
 # Style
 
 Layout/EndAlignment:
@@ -88,7 +91,10 @@ Style/SpecialGlobalVars:
   Enabled: false
 
 Naming/VariableNumber:
-  EnforcedStyle: 'snake_case'
+  EnforcedStyle: "snake_case"
+  AllowedIdentifiers:
+    - sha256
+    - capture3
 
 Naming/MemoizedInstanceVariableName:
   Enabled: false

Gemfile

@@ -29,5 +29,6 @@ group :linting do
 end
 
 group :test do
+  gem "gem_server_conformance", "~> 0.1.5"
   gem "mock_redis"
 end

gemstash.gemspec

@@ -32,13 +32,14 @@ you push your own private gems as well."
   spec.required_ruby_version = ">= 3.1"
 
   spec.add_runtime_dependency "activesupport", ">= 4.2", "< 8"
+  spec.add_runtime_dependency "compact_index", "~> 0.15.0"
   spec.add_runtime_dependency "dalli", ">= 3.2.3", "< 4"
   spec.add_runtime_dependency "faraday", ">= 1", "< 3"
   spec.add_runtime_dependency "faraday_middleware", "~> 1.0"
   spec.add_runtime_dependency "lru_redux", "~> 1.1"
   spec.add_runtime_dependency "psych", ">= 3.2.1"
   spec.add_runtime_dependency "puma", "~> 6.1"
-  spec.add_runtime_dependency "sequel", "~> 5.0"
+  spec.add_runtime_dependency "sequel", "~> 5.85"
   spec.add_runtime_dependency "server_health_check-rack", "~> 0.1"
   spec.add_runtime_dependency "sinatra", ">= 1.4", "< 5.0"
   spec.add_runtime_dependency "terminal-table", "~> 3.0"
@@ -52,8 +53,9 @@ you push your own private gems as well."
   # spec.add_runtime_dependency "mysql2", "~> 0.4"
 
   if RUBY_PLATFORM == "java"
-    spec.add_runtime_dependency "jdbc-sqlite3", "~> 3.8"
+    spec.add_runtime_dependency "jdbc-sqlite3", "~> 3.46"
   else
-    spec.add_runtime_dependency "sqlite3", ">= 1.3", "< 3.0"
+    # SQLite 3.44+ is required for string_agg support
+    spec.add_runtime_dependency "sqlite3", ">= 1.68", "< 3.0"
   end
 end

lib/gemstash.rb

@@ -7,6 +7,7 @@ module Gemstash
   autoload :DB,                  "gemstash/db"
   autoload :Cache,               "gemstash/cache"
   autoload :CLI,                 "gemstash/cli"
+  autoload :CompactIndexBuilder, "gemstash/compact_index_builder"
   autoload :Configuration,       "gemstash/configuration"
   autoload :Dependencies,        "gemstash/dependencies"
   autoload :Env,                 "gemstash/env"

lib/gemstash/cache.rb

@@ -43,7 +43,10 @@ def set_dependency(scope, gem, value)
 
     def invalidate_gem(scope, gem)
       @client.delete("deps/v1/#{scope}/#{gem}")
-      Gemstash::SpecsBuilder.invalidate_stored if scope == "private"
+      if scope == "private"
+        Gemstash::SpecsBuilder.invalidate_stored
+        Gemstash::CompactIndexBuilder.invalidate_stored(gem)
+      end
     end
   end
 

lib/gemstash/compact_index_builder.rb

@@ -0,0 +1,238 @@
+# frozen_string_literal: true
+
+require "active_support/core_ext/string/filters"
+require "compact_index"
+require "gemstash"
+require "stringio"
+require "zlib"
+
+module Gemstash
+  class CompactIndexBuilder
+    include Gemstash::Env::Helper
+    attr_reader :result
+
+    def self.serve(app, ...)
+      app.content_type "text/plain; charset=utf-8"
+      body = new(app.auth, ...).serve
+      app.etag Digest::MD5.hexdigest(body)
+      sha256 = Digest::SHA256.base64digest(body)
+      app.headers "Accept-Ranges" => "bytes", "Digest" => "sha-256=#{sha256}", "Repr-Digest" => "sha-256=:#{sha256}:",
+                  "Content-Length" => body.bytesize.to_s
+      body
+    end
+
+    def self.invalidate_stored(name)
+      storage = Gemstash::Storage.for("private").for("compact_index")
+      storage.resource("names").delete(:names)
+      storage.resource("versions").delete(:versions)
+      storage.resource("info/#{name}").delete(:info)
+    end
+
+    def initialize(auth)
+      @auth = auth
+    end
+
+    def serve
+      check_auth if gemstash_env.config[:protected_fetch]
+      fetch_from_storage
+      return result if result
+
+      build_result
+      store_result
+      result
+    end
+
+  private
+
+    def storage
+      @storage ||= Gemstash::Storage.for("private").for("compact_index")
+    end
+
+    def fetch_from_storage
+      resource = fetch_resource
+      return unless resource.exist?(key)
+
+      @result = resource.load(key).content(key)
+    rescue StandardError
+      # On the off-chance of a race condition between specs.exist? and specs.load
+      @result = nil
+    end
+
+    def store_result
+      fetch_resource.save(key => @result)
+    end
+
+    def check_auth
+      @auth.check("fetch")
+    end
+
+    class Versions < CompactIndexBuilder
+      def fetch_resource
+        storage.resource("versions")
+      end
+
+      def build_result(force_rebuild: false)
+        resource = fetch_resource
+        base = !force_rebuild && resource.exist?("versions.list") && resource.content("versions.list")
+        Tempfile.create("versions.list") do |file|
+          versions_file = CompactIndex::VersionsFile.new(file.path)
+          if base
+            file.write(base)
+            file.close
+            @result = versions_file.contents(
+              compact_index_versions(versions_file.updated_at.to_time)
+            )
+          else
+            ts = Time.now.iso8601
+            versions_file.create(
+              compact_index_public_versions(ts), ts.to_s.sub(/\+00:00\z/, "Z")
+            )
+            @result = file.read
+            resource.save("versions.list" => @result)
+          end
+        end
+      end
+
+    private
+
+      def compact_index_versions(date)
+        all_versions = Sequel::Model.db[<<~SQL.squish, date, date].to_a
+          SELECT r.name as name, v.created_at as date, v.info_checksum as info_checksum, v.number as number, v.platform as platform
+          FROM rubygems AS r, versions AS v
+          WHERE v.rubygem_id = r.id AND
+                v.created_at > ?
+
+          UNION ALL
+
+          SELECT r.name as name, v.yanked_at as date, v.yanked_info_checksum as info_checksum, concat('-', v.number) as number, v.platform as platform
+          FROM rubygems AS r, versions AS v
+          WHERE v.rubygem_id = r.id AND
+                v.indexed is false AND
+                v.yanked_at > ?
+
+          ORDER BY date, number, platform, name
+        SQL
+
+        map_gem_versions(all_versions.map {|v| [v[:name], [v]] })
+      end
+
+      def compact_index_public_versions(date)
+        all_versions = Sequel::Model.db[<<~SQL.squish, date, date].to_a
+          SELECT r.name, v.indexed, COALESCE(v.yanked_at, v.created_at) as stamp,
+                 COALESCE(v.yanked_info_checksum, v.info_checksum) as info_checksum, v.number, v.platform
+          FROM rubygems AS r, versions AS v
+            WHERE v.rubygem_id = r.id AND
+                  (v.created_at <= ? OR v.yanked_at <= ?)
+          ORDER BY name, COALESCE(v.yanked_at, v.created_at), number, platform
+        SQL
+
+        versions_by_gem = all_versions.group_by {|row| row[:name] }
+        versions_by_gem.each_value do |versions|
+          info_checksum = versions.last[:info_checksum]
+          versions.select! {|v| v[:indexed] == true }
+          # Set all versions' info_checksum to work around https://github.com/bundler/compact_index/pull/20
+          versions.each {|v| v[:info_checksum] = info_checksum }
+        end
+
+        map_gem_versions(versions_by_gem)
+      end
+
+      def map_gem_versions(versions_by_gem)
+        versions_by_gem.map do |name, versions|
+          CompactIndex::Gem.new(
+            name,
+            versions.map do |row|
+              CompactIndex::GemVersion.new(
+                row[:number],
+                row[:platform],
+                nil, # sha256
+                row[:info_checksum],
+                nil, # dependencies
+                nil, # version.required_ruby_version,
+                nil, # version.required_rubygems_version
+              )
+            end
+          )
+        end
+      end
+
+      def key
+        :versions
+      end
+    end
+
+    class Info < CompactIndexBuilder
+      def initialize(auth, name)
+        super(auth)
+        @name = name
+      end
+
+      def fetch_resource
+        storage.resource("info/#{@name}")
+      end
+
+      def build_result
+        @result = CompactIndex.info(requirements_and_dependencies)
+      end
+
+    private
+
+      def requirements_and_dependencies
+        DB::Rubygem.association_left_join(versions: :dependencies).
+          where(name: @name).
+          where { versions[:indexed] }.
+          order { [versions[:created_at], versions[:number], versions[:platform], dep_name_agg] }.
+          select_group do
+            [versions[:number], versions[:platform], versions[:sha256], versions[:info_checksum], versions[:required_ruby_version], versions[:required_rubygems_version], versions[:created_at]]
+          end. # rubocop:disable Style/MultilineBlockChain
+          select_more do
+            [coalesce(Sequel.string_agg(dependencies[:requirements], "@").order(dependencies[:rubygem_name], dependencies[:id]), "").as(:dep_req_agg),
+             coalesce(Sequel.string_agg(dependencies[:rubygem_name], ",").order(dependencies[:rubygem_name]), "").as(:dep_name_agg)]
+          end. # rubocop:disable Style/MultilineBlockChain
+          map do |row|
+          reqs = row[:dep_req_agg].split("@")
+          dep_names = row[:dep_name_agg].split(",")
+
+          raise "Dependencies and requirements are not the same size:\n  reqs: #{reqs.inspect}\n  dep_names: #{dep_names.inspect}\n  row: #{row.inspect}" if dep_names.size != reqs.size
+
+          deps = dep_names.zip(reqs).map! do |name, req|
+            CompactIndex::Dependency.new(name, req)
+          end
+
+          CompactIndex::GemVersion.new(
+            row[:number],
+            row[:platform],
+            row[:sha256],
+            nil, # info_checksum
+            deps,
+            row[:required_ruby_version],
+            row[:required_rubygems_version]
+          )
+        end
+      end
+
+      def key
+        :info
+      end
+    end
+
+    class Names < CompactIndexBuilder
+      def fetch_resource
+        storage.resource("names")
+      end
+
+      def build_result
+        names = DB::Rubygem.association_join(:versions).
+                where { versions[:indexed] }.
+                order(:name).group(:name).select_map(:name)
+        @result = CompactIndex.names(names).encode("UTF-8")
+      end
+
+    private
+
+      def key
+        :names
+      end
+    end
+  end
+end

lib/gemstash/db.rb

@@ -10,6 +10,9 @@ module DB
     Sequel::Model.db = Gemstash::Env.current.db
     Sequel::Model.raise_on_save_failure = true
     Sequel::Model.plugin :timestamps, update_on_create: true
+    Sequel::Model.db.extension :error_sql
+    Sequel::Model.db.extension :string_agg
+    Sequel::Model.db.extension :schema_dumper
     autoload :Authorization, "gemstash/db/authorization"
     autoload :CachedRubygem, "gemstash/db/cached_rubygem"
     autoload :Dependency,    "gemstash/db/dependency"

lib/gemstash/db/rubygem.rb

@@ -6,6 +6,8 @@ module Gemstash
   module DB
     # Sequel model for rubygems table.
     class Rubygem < Sequel::Model
+      one_to_many :versions
+
       def self.find_or_insert(spec)
         record = self[name: spec.name]
         return record.id if record