Allow Follower Reads on pg_type Queries (#192)

* Add an AS OF SYSTEM TIME clause to queries on the pg_type catalog.
* Add use_follower_reads database config option
* Reset connection after follower_reads test
* Change use_follower_reads setting to user_follower_reads_for_type_introspection and add Configuration section to README.

Author: keithdoggett

README.md

@@ -22,6 +22,12 @@ development:
   user: <username>
 ```
 
+## Configuration
+
+In addition to the standard adapter settings, CockroachDB also supports the following:
+
+- `use_follower_reads_for_type_introspection`: Use follower reads on queries to the `pg_type` catalog when set to `true`. This helps to speed up initialization by reading historical data, but may not find recently created user-defined types.
+
 ## Working with Spatial Data
 
 The adapter uses [RGeo](https://github.com/rgeo/rgeo) and [RGeo-ActiveRecord](https://github.com/rgeo/rgeo-activerecord) to represent geometric and geographic data as Ruby objects and easily interface them with the adapter. The following is a brief introduction to RGeo and tips to help setup your spatial application. More documentation about RGeo can be found in the [YARD Docs](https://rubydoc.info/github/rgeo/rgeo) and [wiki](https://github.com/rgeo/rgeo/wiki).

lib/active_record/connection_adapters/cockroachdb_adapter.rb

@@ -403,6 +403,100 @@ def extract_empty_array_from_default(default)
           return "{}"
         end
 
+        # override
+        # This method loads info about data types from the database to
+        # populate the TypeMap.
+        #
+        # Currently, querying from the pg_type catalog can be slow due to geo-partitioning
+        # so this modified query uses AS OF SYSTEM TIME '-10s' to read historical data.
+        def load_additional_types(oids = nil)
+          if @config[:use_follower_reads_for_type_introspection]
+            initializer = OID::TypeMapInitializer.new(type_map)
+
+            query = <<~SQL
+              SELECT t.oid, t.typname, t.typelem, t.typdelim, t.typinput, r.rngsubtype, t.typtype, t.typbasetype
+              FROM pg_type as t
+              LEFT JOIN pg_range as r ON oid = rngtypid AS OF SYSTEM TIME '-10s'
+            SQL
+
+            if oids
+              query += "WHERE t.oid IN (%s)" % oids.join(", ")
+            else
+              query += initializer.query_conditions_for_initial_load
+            end
+
+            execute_and_clear(query, "SCHEMA", []) do |records|
+              initializer.run(records)
+            end
+          else
+            super
+          end
+        rescue ActiveRecord::StatementInvalid => e
+          raise e unless e.cause.is_a? PG::InvalidCatalogName
+          # use original if database is younger than 10s
+          super
+        end
+
+        # override
+        # This method maps data types to their proper decoder.
+        #
+        # Currently, querying from the pg_type catalog can be slow due to geo-partitioning
+        # so this modified query uses AS OF SYSTEM TIME '-10s' to read historical data.
+        def add_pg_decoders
+          if @config[:use_follower_reads_for_type_introspection]
+            @default_timezone = nil
+            @timestamp_decoder = nil
+
+            coders_by_name = {
+              "int2" => PG::TextDecoder::Integer,
+              "int4" => PG::TextDecoder::Integer,
+              "int8" => PG::TextDecoder::Integer,
+              "oid" => PG::TextDecoder::Integer,
+              "float4" => PG::TextDecoder::Float,
+              "float8" => PG::TextDecoder::Float,
+              "numeric" => PG::TextDecoder::Numeric,
+              "bool" => PG::TextDecoder::Boolean,
+              "timestamp" => PG::TextDecoder::TimestampUtc,
+              "timestamptz" => PG::TextDecoder::TimestampWithTimeZone,
+            }
+
+            known_coder_types = coders_by_name.keys.map { |n| quote(n) }
+            query = <<~SQL % known_coder_types.join(", ")
+              SELECT t.oid, t.typname
+              FROM pg_type as t AS OF SYSTEM TIME '-10s'
+              WHERE t.typname IN (%s)
+            SQL
+
+            coders = execute_and_clear(query, "SCHEMA", []) do |result|
+              result
+                .map { |row| construct_coder(row, coders_by_name[row["typname"]]) }
+                .compact
+            end
+
+            map = PG::TypeMapByOid.new
+            coders.each { |coder| map.add_coder(coder) }
+            @connection.type_map_for_results = map
+
+            @type_map_for_results = PG::TypeMapByOid.new
+            @type_map_for_results.default_type_map = map
+            @type_map_for_results.add_coder(PG::TextDecoder::Bytea.new(oid: 17, name: "bytea"))
+
+            # extract timestamp decoder for use in update_typemap_for_default_timezone
+            @timestamp_decoder = coders.find { |coder| coder.name == "timestamp" }
+            update_typemap_for_default_timezone
+          else
+            super
+          end
+        rescue ActiveRecord::StatementInvalid => e
+          raise e unless e.cause.is_a? PG::InvalidCatalogName
+          # use original if database is younger than 10s
+          super
+        end
+
+        def arel_visitor
+          Arel::Visitors::CockroachDB.new(self)
+        end
+
       # end private
     end
   end

test/cases/adapters/postgresql/postgresql_adapter_test.rb

@@ -15,13 +15,34 @@ def setup
         @connection_handler = ActiveRecord::Base.connection_handler
       end
 
+      def teardown
+        # use connection without follower_reads
+        database_config = { "adapter" => "cockroachdb", "database" => "activerecord_unittest" }
+        ar_config = ActiveRecord::Base.configurations["arunit"]
+        database_config.update(ar_config)
+
+        ActiveRecord::Base.establish_connection(database_config)
+      end
+
       def test_database_exists_returns_false_when_the_database_does_not_exist
         config = ActiveRecord::Base.configurations["arunit"]
         bad_config = config.dup
         bad_config[:database] = "non_extant_database"
         assert_not ActiveRecord::ConnectionAdapters::CockroachDBAdapter.database_exists?(bad_config),
           "expected database #{bad_config[:database]} to not exist"
       end
+
+      def test_using_follower_reads_connects_properly
+        database_config = { "use_follower_reads_for_type_introspection": true, "adapter" => "cockroachdb", "database" => "activerecord_unittest" }
+        ar_config = ActiveRecord::Base.configurations["arunit"]
+        database_config.update(ar_config)
+
+        ActiveRecord::Base.establish_connection(database_config)
+        conn = ActiveRecord::Base.connection
+        conn_config = conn.instance_variable_get("@config")
+
+        assert conn_config[:use_follower_reads_for_type_introspection]
+      end
     end
   end
 end