Merge pull request #35256 from sponomarev/feature/cable-postgres-identity

Add ActionCable subscription connection identifier to PostgreSQL adapter

actioncable/CHANGELOG.md

+*   Add PostgreSQL subscription connection identificator.
+
+    Now you can distinguish Action Cable PostgreSQL subscription connections among others.
+    Also, you can set custom `id` in `cable.yml` configuration.
+
+    ```sql
+    SELECT application_name FROM pg_stat_activity;
+    /*
+        application_name
+    ------------------------
+    psql
+    ActionCable-PID-42
+    (2 rows)
+    */
+     ```
+
+     *Sergey Ponomarev*
+
 *   Subscription confirmations and rejections are now logged at the `DEBUG` level instead of `INFO`.
 
     *DHH*
 
-
 Please check [6-0-stable](https://github.com/rails/rails/blob/6-0-stable/actioncable/CHANGELOG.md) for previous changes.

actioncable/lib/action_cable/subscription_adapter/base.rb

@@ -25,6 +25,10 @@ def unsubscribe(channel, message_callback)
       def shutdown
         raise NotImplementedError
       end
+
+      def identifier
+        @server.config.cable[:id] ||= "ActionCable-PID-#{$$}"
+      end
     end
   end
 end

actioncable/lib/action_cable/subscription_adapter/postgresql.rb

@@ -42,6 +42,7 @@ def with_subscriptions_connection(&block) # :nodoc:
         pg_conn = ar_conn.raw_connection
 
         verify!(pg_conn)
+        pg_conn.exec("SET application_name = #{pg_conn.escape_identifier(identifier)}")
         yield pg_conn
       ensure
         ar_conn.disconnect!

actioncable/lib/action_cable/subscription_adapter/redis.rb

@@ -15,7 +15,6 @@ class Redis < Base # :nodoc:
       # Overwrite this factory method for Redis connections if you want to use a different Redis library than the redis gem.
       # This is needed, for example, when using Makara proxies for distributed Redis.
       cattr_accessor :redis_connector, default: ->(config) do
-        config[:id] ||= "ActionCable-PID-#{$$}"
         ::Redis.new(config.except(:adapter, :channel_prefix))
       end
 
@@ -57,7 +56,7 @@ def redis_connection_for_broadcasts
         end
 
         def redis_connection
-          self.class.redis_connector.call(@server.config.cable)
+          self.class.redis_connector.call(@server.config.cable.merge(id: identifier))
         end
 
         class Listener < SubscriberMap

actioncable/test/subscription_adapter/postgresql_test.rb

@@ -64,4 +64,24 @@ def active?
 
     assert adapter.active?
   end
+
+  def test_default_subscription_connection_identifier
+    subscribe_as_queue("channel") { }
+
+    identifiers = ActiveRecord::Base.connection.exec_query("SELECT application_name FROM pg_stat_activity").rows
+    assert_includes identifiers, ["ActionCable-PID-#{$$}"]
+  end
+
+  def test_custom_subscription_connection_identifier
+    server = ActionCable::Server::Base.new
+    server.config.cable = cable_config.merge(id: "hello-world-42").with_indifferent_access
+    server.config.logger = Logger.new(StringIO.new).tap { |l| l.level = Logger::UNKNOWN }
+
+    adapter = server.config.pubsub_adapter.new(server)
+
+    subscribe_as_queue("channel", adapter) { }
+
+    identifiers = ActiveRecord::Base.connection.exec_query("SELECT application_name FROM pg_stat_activity").rows
+    assert_includes identifiers, ["hello-world-42"]
+  end
 end

actioncable/test/subscription_adapter/redis_test.rb

@@ -33,24 +33,50 @@ def cable_config
   end
 end
 
-class RedisAdapterTest::Connector < ActionCable::TestCase
-  test "excludes adapter and channel prefix" do
-    config = { url: 1, host: 2, port: 3, db: 4, password: 5, id: "Some custom ID" }
+class RedisAdapterTest::ConnectorDefaultID < ActionCable::TestCase
+  def setup
+    server = ActionCable::Server::Base.new
+    server.config.cable = cable_config.merge(adapter: "redis").with_indifferent_access
+    server.config.logger = Logger.new(StringIO.new).tap { |l| l.level = Logger::UNKNOWN }
 
-    assert_called_with ::Redis, :new, [ config ] do
-      connect config.merge(adapter: "redis", channel_prefix: "custom")
+    @adapter = server.config.pubsub_adapter.new(server)
   end
+
+  def cable_config
+    { url: 1, host: 2, port: 3, db: 4, password: 5 }
+  end
+
+  def connection_id
+    "ActionCable-PID-#{$$}"
   end
 
-  test "adds default id if it is not specified" do
-    config = { url: 1, host: 2, port: 3, db: 4, password: 5, id: "ActionCable-PID-#{$$}" }
+  def expected_connection
+    cable_config.merge(id: connection_id)
+  end
 
-    assert_called_with ::Redis, :new, [ config ] do
-      connect config.except(:id)
+  test "sets connection id for connection" do
+    assert_called_with ::Redis, :new, [ expected_connection.stringify_keys ] do
+      @adapter.send(:redis_connection)
+    end
   end
+end
+
+class RedisAdapterTest::ConnectorCustomID < RedisAdapterTest::ConnectorDefaultID
+  def cable_config
+    super.merge(id: connection_id)
+  end
+
+  def connection_id
+    "Some custom ID"
+  end
+end
+
+class RedisAdapterTest::ConnectorWithExcluded < RedisAdapterTest::ConnectorDefaultID
+  def cable_config
+    super.merge(adapter: "redis", channel_prefix: "custom")
   end
 
-  def connect(config)
-    ActionCable::SubscriptionAdapter::Redis.redis_connector.call(config)
+  def expected_connection
+    super.except(:adapter, :channel_prefix)
   end
 end