Add pg_schema_caching extension, for reloading OIDs for custom types …

…when loading cached schema

While PostgreSQL uses the same OIDs for all built-in types, custom
types have OIDs that differ for each database, even for databases
built from the same migrations.  That means that schema caches
built from development or test databases, if loaded into a
Database object for the production database, will have incorrect
oids for custom types.

This avoids the issue by replacing custom oids with :custom in
dumped schema. When loading schema, a single query is done to get
the oids for each custom type in the dumped schema, and the column
schema hashes are then updated to set the correct oid.

CHANGELOG

@@ -1,5 +1,7 @@
 === master
 
+* Add pg_schema_caching extension, for reloading OIDs for custom types when loading cached schema (jeremyevans)
+
 * Make Database#schema hashes include :comment field on MySQL and PostgreSQL (Bahanix) (#2248, #2249)
 
 * Add inspect_pk plugin to make it easier to retrieve model instance based on inspect output (jeremyevans)

lib/sequel/extensions/pg_schema_caching.rb

@@ -0,0 +1,90 @@
+# frozen-string-literal: true
+#
+# The pg_schema_caching extension builds on top of the schema_caching
+# extension, and allows it to handle custom PostgreSQL types. On
+# PostgreSQL, column schema hashes include an :oid entry for the OID
+# for the column's type.  For custom types, this OID is dependent on
+# the PostgreSQL database, so in most cases, test and development
+# versions of the same database, created with the same migrations,
+# will have different OIDs.
+#
+# To fix this case, the pg_schema_caching extension removes custom
+# OIDs from the schema cache when dumping the schema, replacing them
+# with a placeholder. When loading the cached schema, the Database
+# object makes a single query to get the OIDs for all custom types
+# used by the cached schema, and it updates all related column
+# schema hashes to set the correct :oid entry for the current
+# database.
+#
+# Related module: Sequel::Postgres::SchemaCaching
+
+require_relative "schema_caching"
+
+module Sequel
+  module Postgres
+    module SchemaCaching
+      include Sequel::SchemaCaching
+
+      private
+
+      # Load custom oids from database when loading schema cache file.
+      def load_schema_cache_file(file)
+        set_custom_oids_for_cached_schema(super)
+      end
+
+      # Find all column schema hashes that use custom types.
+      # Load the oids for custom types in a single query, and update
+      # each related column schema hash with the correct oid.
+      def set_custom_oids_for_cached_schema(schemas)
+        custom_oid_rows = {}
+
+        schemas.each_value do |cols|
+          cols.each do |_, h|
+            if h[:oid] == :custom
+              (custom_oid_rows[h[:db_type]] ||= []) << h
+            end
+          end
+        end
+
+        unless custom_oid_rows.empty?
+          from(:pg_type).where(:typname=>custom_oid_rows.keys).select_hash(:typname, :oid).each do |name, oid|
+            custom_oid_rows.delete(name).each do |row|
+              row[:oid] = oid
+            end
+          end
+        end
+
+        unless custom_oid_rows.empty?
+          warn "Could not load OIDs for the following custom types: #{custom_oid_rows.keys.sort.join(", ")}", uplevel: 3
+
+          schemas.keys.each do |k|
+            if schemas[k].any?{|_,h| h[:oid] == :custom}
+              # Remove schema entry for table, so it will be queried at runtime to get the correct oids
+              schemas.delete(k)
+            end
+          end
+        end
+
+        schemas
+      end
+
+      # Replace :oid entries for custom types with :custom.
+      def dumpable_schema_cache
+        sch = super
+
+        sch.each_value do |cols|
+          cols.each do |_, h|
+            if (oid = h[:oid]) && oid >= 10000
+              h[:oid] = :custom
+            end
+          end
+        end
+
+        sch
+      end
+    end
+  end
+
+  Database.register_extension(:pg_schema_caching, Postgres::SchemaCaching)
+end
+

lib/sequel/extensions/schema_caching.rb

@@ -51,14 +51,7 @@ module Sequel
   module SchemaCaching
     # Dump the cached schema to the filename given in Marshal format.
     def dump_schema_cache(file)
-      sch = {}
-      @schemas.sort.each do |k,v|
-        sch[k] = v.map do |c, h|
-          h = Hash[h]
-          h.delete(:callable_default)
-          [c, h]
-        end
-      end
+      sch = dumpable_schema_cache
       File.open(file, 'wb'){|f| f.write(Marshal.dump(sch))}
       nil
     end
@@ -72,7 +65,7 @@ def dump_schema_cache?(file)
     # Replace the schema cache with the data from the given file, which
     # should be in Marshal format.
     def load_schema_cache(file)
-      @schemas = Marshal.load(File.read(file))
+      @schemas = load_schema_cache_file(file)
       @schemas.each_value{|v| schema_post_process(v)}
       nil
     end
@@ -82,6 +75,28 @@ def load_schema_cache(file)
     def load_schema_cache?(file)
       load_schema_cache(file) if File.exist?(file)
     end
+
+    private
+
+    # Return the deserialized schema cache file.
+    def load_schema_cache_file(file)
+      Marshal.load(File.read(file))
+    end
+
+    # A dumpable version of the schema cache.
+    def dumpable_schema_cache
+      sch = {}
+
+      @schemas.sort.each do |k,v|
+        sch[k] = v.map do |c, h|
+          h = Hash[h]
+          h.delete(:callable_default)
+          [c, h]
+        end
+      end
+
+      sch
+    end
   end
 
   Database.register_extension(:schema_caching, SchemaCaching)

spec/adapters/postgres_spec.rb

@@ -6355,3 +6355,37 @@ def check(ds)
     @m1.order(:i1).all.must_equal [{:i1=>1, :a=>100}, {:i1=>3, :a=>40}, {:i1=>4, :a=>-15}]
   end
 end if DB.server_version >= 170000
+
+describe 'pg_schema_caching_extension' do
+  before do
+    @db = DB
+    @cache_file = "spec/files/pg_schema_caching-spec-#{$$}.cache"
+  end
+  after do
+    @db.drop_table?(:test_pg_schema_caching, :test_pg_schema_caching_type)
+    File.delete(@cache_file) if File.file?(@cache_file)
+  end
+
+  it "should encode custom type OIDs as :custom, and reload them" do
+    @db.create_table(:test_pg_schema_caching_type) do
+      Integer :id
+      String :name
+    end
+    @db.create_table(:test_pg_schema_caching) do
+      test_pg_schema_caching_type :t
+    end
+
+    oid = @db.schema(:test_pg_schema_caching)[0][1][:oid]
+    oid.must_be_kind_of Integer
+
+    @db.extension :pg_schema_caching
+    @db.dump_schema_cache(@cache_file)
+
+    @db.instance_variable_get(:@schemas).delete('"test_pg_schema_caching"')
+    cache = Marshal.load(File.binread(@cache_file))
+    cache['"test_pg_schema_caching"'][0][1][:oid].must_equal :custom
+
+    @db.load_schema_cache(@cache_file)
+    @db.schema(:test_pg_schema_caching)[0][1][:oid].must_equal oid
+  end
+end

spec/extensions/pg_schema_caching_spec.rb

@@ -0,0 +1,67 @@
+require_relative "spec_helper"
+
+describe "pg_schema_caching extension" do
+  before do
+    @db = Sequel.connect('mock://postgres').extension(:pg_schema_caching)
+    @schemas = {
+      '"table1"'=>[
+        [:column1, {:oid=>11111, :db_type=>"custom_type", :default=>"nextval('table_id_seq'::regclass)", :allow_null=>false, :primary_key=>true, :type=>:integer, :ruby_default=>nil}],
+        [:column2, {:oid=>1111, :db_type=>"integer", :default=>"nextval('table_id_seq'::regclass)", :allow_null=>false, :primary_key=>true, :type=>:integer, :ruby_default=>nil}],
+      ],
+      '"table2"'=>[
+        [:column3, {:oid=>1111, :db_type=>"integer", :default=>"nextval('table_id_seq'::regclass)", :allow_null=>false, :primary_key=>true, :type=>:integer, :ruby_default=>nil}],
+      ],
+      '"table3"'=>[
+        [:column4, {:oid=>11112, :db_type=>"custom_type2", :default=>"nextval('table_id_seq'::regclass)", :allow_null=>false, :primary_key=>true, :type=>:integer, :ruby_default=>nil}],
+        [:column5, {:oid=>1111, :db_type=>"integer", :default=>"nextval('table_id_seq'::regclass)", :allow_null=>false, :primary_key=>true, :type=>:integer, :ruby_default=>nil}],
+      ]
+    }
+    @filename = "spec/files/test_schema_#$$.dump" 
+    @db.instance_variable_set(:@schemas, @schemas)
+  end
+  after do
+    File.delete(@filename) if File.exist?(@filename)
+  end
+
+  it "Database#dump_schema_cache should dump cached schema to the given file without custom oids" do
+    File.exist?(@filename).must_equal false
+    @db.dump_schema_cache(@filename)
+    File.exist?(@filename).must_equal true
+    cache = Marshal.load(File.binread(@filename))
+    cache['"table1"'][0][1][:oid].must_equal :custom
+    cache['"table1"'][1][1][:oid].must_equal 1111
+    cache['"table2"'][0][1][:oid].must_equal 1111
+    cache['"table3"'][0][1][:oid].must_equal :custom
+    cache['"table3"'][1][1][:oid].must_equal 1111
+  end
+
+  it "Database#load_schema_cache should load cached schema, using a single query for custom type oids" do
+    @db.dump_schema_cache(@filename)
+    @db.fetch = [{:typname=>"custom_type2", :oid=>22221}, {:typname=>"custom_type", :oid=>22222}]
+    @db.load_schema_cache(@filename)
+    @db.schema(:table1)[0][1][:oid].must_equal 22222
+    @db.schema(:table1)[1][1][:oid].must_equal 1111
+    @db.schema(:table2)[0][1][:oid].must_equal 1111
+    @db.schema(:table3)[0][1][:oid].must_equal 22221
+    @db.schema(:table3)[1][1][:oid].must_equal 1111
+    @db.sqls.must_equal ["SELECT \"typname\", \"oid\" FROM \"pg_type\" WHERE (\"typname\" IN ('custom_type', 'custom_type2'))"]
+  end
+
+  it "Database#load_schema_cache should load cached schema without issuing a query if there are no custom type oids" do
+    @schemas.delete('"table1"')
+    @schemas.delete('"table3"')
+    @db.dump_schema_cache(@filename)
+    @db.load_schema_cache(@filename)
+    @db.sqls.must_equal []
+  end
+
+  it "Database#load_schema_cache should warn if custom type oids present in cache are not found in the database, and remove schema entry from cache" do
+    @db.dump_schema_cache(@filename)
+    @db.fetch = [{:typname=>"custom_type2", :oid=>22221}]
+    a = []
+    @db.define_singleton_method(:warn){|*args| a.replace(args)}
+    @db.load_schema_cache(@filename)
+    a.must_equal ["Could not load OIDs for the following custom types: custom_type", {:uplevel=>3}]
+    @db.instance_variable_get(:@schemas).keys.must_equal(%w'"table2" "table3"')
+  end
+end

www/pages/plugins.html.erb

@@ -739,6 +739,10 @@
 <span class="ul__span">Adds support for PostgreSQL row-valued/composite types.</span>
 </li>
 <li class="ul__li ul__li--grid">
+<a class="a" href="rdoc-plugins/files/lib/sequel/extensions/pg_schema_caching_rb.html">pg_schema_caching </a>
+<span class="ul__span">Builds on schema_caching extension, reloading OIDs for custom types when loading cached schema.</span>
+</li>
+<li class="ul__li ul__li--grid">
 <a class="a" href="rdoc-plugins/files/lib/sequel/extensions/pg_static_cache_updater_rb.html">pg_static_cache_updater </a>
 <span class="ul__span">Listens for changes to underlying tables in order to automatically update models using the static_cache plugin.</span>
 </li>