Support gzip for schema cache

Katrina Owen · Katrina Owen · commit ffc3af63e00e · 2020-02-24T14:56:13.000-07:00
This adds gzip support for both the YAML and the Marshal serialization
strategies.

Particularly large schema caches can become a problem when deploying to
Kubernetes, as there is currently a 1*1024*1024 byte limit for the
ConfigMap. For large databases, the schema cache can exceed this limit.
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.include?(".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,7 +155,7 @@ def clear_data_source_cache!(name)
       def dump_to(filename)
         clear!
         connection.data_sources.each { |table| add(table) }
-        File.atomic_write(filename) { |f|
+        open(filename) { |f|
           if filename.include?(".dump")
             f.write(Marshal.dump(self))
           else
@@ -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")
