Merge pull request #196 from keithdoggett/increment-feature-counter

Call Telemetry on Connection Pool Initialization

Author: solongordon

README.md

@@ -27,6 +27,7 @@ development:
 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.
+- `disable_cockroachdb_telemetry`: Determines if a telemetry call is made to the database when the connection pool is initialized. Setting this to `true` will prevent the call from being made.
 
 ## Working with Spatial Data
 

build/config.teamcity.yml

@@ -11,6 +11,7 @@ connections:
       user: root
       requiressl: disable
       min_messages: warning
+      disable_cockroachdb_telemetry: true
     arunit_without_prepared_statements:
       database: activerecord_unittest
       host: localhost
@@ -19,10 +20,12 @@ connections:
       requiressl: disable
       min_messages: warning
       prepared_statements: false
+      disable_cockroachdb_telemetry: true
     arunit2:
       database: activerecord_unittest2
       host: localhost
       port: 26257
       user: root
       requiressl: disable
       min_messages: warning
+      disable_cockroachdb_telemetry: true

lib/active_record/connection_adapters/cockroachdb_adapter.rb

@@ -56,6 +56,30 @@ def cockroachdb_connection(config)
 
 module ActiveRecord
   module ConnectionAdapters
+    module CockroachDBConnectionPool
+      def initialize(pool_config)
+        super(pool_config)
+        disable_telemetry = pool_config.db_config.configuration_hash[:disable_cockroachdb_telemetry]
+        adapter = pool_config.db_config.configuration_hash[:adapter]
+        return if disable_telemetry || adapter != "cockroachdb"
+
+        with_connection do |conn|
+          if conn.active?
+            begin
+              query = "SELECT crdb_internal.increment_feature_counter('ActiveRecord %d.%d')"
+              conn.execute(query % [ActiveRecord::VERSION::MAJOR, ActiveRecord::VERSION::MINOR])
+            rescue ActiveRecord::StatementInvalid
+              # The increment_feature_counter built-in is not supported on this
+              # CockroachDB version. Ignore.
+            rescue StandardError => e
+              conn.logger.warn "Unexpected error when incrementing feature counter: #{e}"
+            end
+          end
+        end
+      end
+    end
+    ConnectionPool.prepend(CockroachDBConnectionPool)
+
     class CockroachDBAdapter < PostgreSQLAdapter
       ADAPTER_NAME = "CockroachDB".freeze
       DEFAULT_PRIMARY_KEY = "rowid"
@@ -194,6 +218,7 @@ def max_identifier_length
 
       def initialize(connection, logger, conn_params, config)
         super(connection, logger, conn_params, config)
+
         crdb_version_string = query_value("SHOW crdb_version")
         if crdb_version_string.include? "v1."
           version_num = 1

test/cases/adapters/postgresql/postgresql_adapter_test.rb

@@ -16,7 +16,7 @@ def setup
       end
 
       def teardown
-        # use connection without follower_reads
+        # use connection without follower_reads and telemetry
         database_config = { "adapter" => "cockroachdb", "database" => "activerecord_unittest" }
         ar_config = ActiveRecord::Base.configurations.configs_for(env_name: "arunit", name: "primary")
         database_config.update(ar_config.configuration_hash)
@@ -38,6 +38,19 @@ def test_database_exists_returns_true_when_the_database_exists
           "expected database #{db_config.database} to exist"
       end
 
+      def test_using_telemetry_builtin_connects_properly
+        database_config = { "adapter" => "cockroachdb", "database" => "activerecord_unittest" }
+        ar_config = ActiveRecord::Base.configurations.configs_for(env_name: "arunit", name: "primary")
+        database_config.update(ar_config.configuration_hash)
+        database_config[:disable_cockroachdb_telemetry] = false
+
+        ActiveRecord::Base.establish_connection(database_config)
+        conn = ActiveRecord::Base.connection
+        conn_config = conn.instance_variable_get("@config")
+
+        assert_equal(false, conn_config[:disable_cockroachdb_telemetry])
+      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.configs_for(env_name: "arunit", name: "primary")

test/config.yml

@@ -5,14 +5,17 @@ connections:
       host: localhost
       port: 26257
       user: root
+      disable_cockroachdb_telemetry: true
     arunit_without_prepared_statements:
       min_messages: warning
       prepared_statements: false
       host: localhost
       port: 26257
       user: root
+      disable_cockroachdb_telemetry: true
     arunit2:
       min_messages: warning
       host: localhost
       port: 26257
       user: root
+      disable_cockroachdb_telemetry: true