Merge pull request rails#38566 from kytrinyx/gzip-schema-cache

rafaelfranca · web-flow · commit 97cd1c8078f3 · 2020-02-24T17:37:22.000-05:00
Support gzip for the schema cache
diff --git a/activerecord/lib/active_record/connection_adapters/schema_cache.rb b/activerecord/lib/active_record/connection_adapters/schema_cache.rb
@@ -8,9 +8,21 @@ class SchemaCache
       def self.load_from(filename)
         return unless File.file?(filename)
 
-        file = File.read(filename)
-        filename.end_with?(".dump") ? Marshal.load(file) : YAML.load(file)
+        read(filename) do |file|
+          filename.include?(".dump") ? Marshal.load(file) : YAML.load(file)
+        end
+      end
+
+      def self.read(filename, &block)
+        if File.extname(filename) == ".gz"
+          Zlib::GzipReader.open(filename) { |gz|
+            yield gz.read
+          }
+        else
+          yield File.read(filename)
+        end
       end
+      private_class_method :read
 
       attr_reader :version
       attr_accessor :connection
@@ -143,8 +155,8 @@ def clear_data_source_cache!(name)
       def dump_to(filename)
         clear!
         connection.data_sources.each { |table| add(table) }
-        File.atomic_write(filename) { |f|
-          if filename.end_with?(".dump")
+        open(filename) { |f|
+          if filename.include?(".dump")
             f.write(Marshal.dump(self))
           else
             f.write(YAML.dump(self))
@@ -194,6 +206,19 @@ def deep_deduplicate(value)
         def prepare_data_sources
           connection.data_sources.each { |source| @data_sources[source] = true }
         end
+
+        def open(filename)
+          File.atomic_write(filename) do |file|
+            if File.extname(filename) == ".gz"
+              zipper = Zlib::GzipWriter.new file
+              yield zipper
+              zipper.flush
+              zipper.close
+            else
+              yield file
+            end
+          end
+        end
     end
   end
 end
diff --git a/activerecord/test/cases/connection_adapters/schema_cache_test.rb b/activerecord/test/cases/connection_adapters/schema_cache_test.rb
@@ -43,6 +43,49 @@ def test_yaml_dump_and_load
         tempfile.unlink
       end
 
+      def test_yaml_dump_and_load_with_gzip
+        # Create an empty cache.
+        cache = SchemaCache.new @connection
+
+        tempfile = Tempfile.new(["schema_cache-", ".yml.gz"])
+        # Dump it. It should get populated before dumping.
+        cache.dump_to(tempfile.path)
+
+        # Unzip and load manually.
+        cache = Zlib::GzipReader.open(tempfile.path) { |gz| YAML.load(gz.read) }
+
+        # Give it a connection. Usually the connection
+        # would get set on the cache when it's retrieved
+        # from the pool.
+        cache.connection = @connection
+
+        assert_no_queries do
+          assert_equal 12, cache.columns("posts").size
+          assert_equal 12, cache.columns_hash("posts").size
+          assert cache.data_sources("posts")
+          assert_equal "id", cache.primary_keys("posts")
+          assert_equal 1, cache.indexes("posts").size
+          assert_equal @database_version.to_s, cache.database_version.to_s
+        end
+
+        # Load the cache the usual way.
+        cache = SchemaCache.load_from(tempfile.path)
+
+        # Give it a connection.
+        cache.connection = @connection
+
+        assert_no_queries do
+          assert_equal 12, cache.columns("posts").size
+          assert_equal 12, cache.columns_hash("posts").size
+          assert cache.data_sources("posts")
+          assert_equal "id", cache.primary_keys("posts")
+          assert_equal 1, cache.indexes("posts").size
+          assert_equal @database_version.to_s, cache.database_version.to_s
+        end
+      ensure
+        tempfile.unlink
+      end
+
       def test_yaml_loads_5_1_dump
         cache = SchemaCache.load_from(schema_dump_path)
         cache.connection = @connection
@@ -162,6 +205,43 @@ def test_marshal_dump_and_load_via_disk
         tempfile.unlink
       end
 
+      def test_marshal_dump_and_load_with_gzip
+        # Create an empty cache.
+        cache = SchemaCache.new @connection
+
+        tempfile = Tempfile.new(["schema_cache-", ".dump.gz"])
+        # Dump it. It should get populated before dumping.
+        cache.dump_to(tempfile.path)
+
+        # Load a new cache manually.
+        cache = Zlib::GzipReader.open(tempfile.path) { |gz| Marshal.load(gz.read) }
+        cache.connection = @connection
+
+        assert_no_queries do
+          assert_equal 12, cache.columns("posts").size
+          assert_equal 12, cache.columns_hash("posts").size
+          assert cache.data_sources("posts")
+          assert_equal "id", cache.primary_keys("posts")
+          assert_equal 1, cache.indexes("posts").size
+          assert_equal @database_version.to_s, cache.database_version.to_s
+        end
+
+        # Load a new cache.
+        cache = SchemaCache.load_from(tempfile.path)
+        cache.connection = @connection
+
+        assert_no_queries do
+          assert_equal 12, cache.columns("posts").size
+          assert_equal 12, cache.columns_hash("posts").size
+          assert cache.data_sources("posts")
+          assert_equal "id", cache.primary_keys("posts")
+          assert_equal 1, cache.indexes("posts").size
+          assert_equal @database_version.to_s, cache.database_version.to_s
+        end
+      ensure
+        tempfile.unlink
+      end
+
       def test_data_source_exist
         assert @cache.data_source_exists?("posts")
         assert_not @cache.data_source_exists?("foo")
