Add multiple database (roles) support

Nikita Bulai · nbulaj · commit 157f7516ac88 · 2026-02-24T21:45:40.000+03:00
diff --git a/.gitignore b/.gitignore
@@ -20,3 +20,4 @@ gemfiles/*.lock
 coverage
 *.gem
 gemfiles/vendor
+vendor/bundle/
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -8,6 +8,7 @@ User-visible changes worth mentioning.
 ## main
 
 Add your entry here.
+- [#1791] Add support for Rails read replicas with automatic role switching via `enable_multiple_databases` configuration option
 - [#1792] Consider expires_in when clear expired tokens with StaleRecordsCleaner.
 - [#1790] Fix race condition in refresh token revocation check by moving InvalidGrantReuse check inside the lock block
 - [#1788] Fix regex for basic auth to be case-insensitive
diff --git a/lib/doorkeeper.rb b/lib/doorkeeper.rb
@@ -101,6 +101,10 @@ module Models
     autoload :ResourceOwnerable, "doorkeeper/models/concerns/resource_ownerable"
     autoload :Revocable, "doorkeeper/models/concerns/revocable"
     autoload :SecretStorable, "doorkeeper/models/concerns/secret_storable"
+
+    module Concerns
+      autoload :WriteToPrimary, "doorkeeper/models/concerns/write_to_primary"
+    end
   end
 
   module Orm
diff --git a/lib/doorkeeper/config.rb b/lib/doorkeeper/config.rb
@@ -99,6 +99,22 @@ def reuse_access_token
         @config.instance_variable_set(:@reuse_access_token, true)
       end
 
+      # Enable support for multiple database configurations with read replicas.
+      # When enabled, wraps database write operations to ensure they use the primary
+      # (writable) database when automatic role switching is enabled.
+      #
+      # For ActiveRecord (Rails 6.1+), this uses `ActiveRecord::Base.connected_to(role: :writing)`.
+      # Other ORM extensions can implement their own primary database targeting logic.
+      #
+      # This prevents `ActiveRecord::ReadOnlyError` when using read replicas with Rails
+      # automatic role switching. Enable this if your application uses multiple databases
+      # with automatic role switching for read replicas.
+      #
+      # See: https://guides.rubyonrails.org/active_record_multiple_databases.html#activating-automatic-role-switching
+      def enable_multiple_databases
+        @config.instance_variable_set(:@enable_multiple_databases, true)
+      end
+
       # Choose to use the url path for native autorization codes 
       # Enabling this flag sets the authorization code response route for
       # native redirect uris to oauth/authorize/<code>. The default is
@@ -437,6 +453,7 @@ def configure_secrets_for(type, using:, fallback:)
            end)
 
     attr_reader :reuse_access_token,
+                :enable_multiple_databases,
                 :token_secret_fallback_strategy,
                 :application_secret_fallback_strategy
 
diff --git a/lib/doorkeeper/models/access_grant_mixin.rb b/lib/doorkeeper/models/access_grant_mixin.rb
@@ -12,6 +12,7 @@ module AccessGrantMixin
     include Models::SecretStorable
     include Models::Scopes
     include Models::ResourceOwnerable
+    include Models::Concerns::WriteToPrimary
 
     # Never uses PKCE if PKCE migrations were not generated
     def uses_pkce?
@@ -40,12 +41,14 @@ def by_token(token)
       #   instance of the Resource Owner model or it's ID
       #
       def revoke_all_for(application_id, resource_owner, clock = Time)
-        by_resource_owner(resource_owner)
-          .where(
-            application_id: application_id,
-            revoked_at: nil,
-          )
-          .update_all(revoked_at: clock.now.utc)
+        with_primary_role do
+          by_resource_owner(resource_owner)
+            .where(
+              application_id: application_id,
+              revoked_at: nil,
+            )
+            .update_all(revoked_at: clock.now.utc)
+        end
       end
 
       # Implements PKCE code_challenge encoding without base64 padding as described in the spec.
diff --git a/lib/doorkeeper/models/access_token_mixin.rb b/lib/doorkeeper/models/access_token_mixin.rb
@@ -14,6 +14,7 @@ module AccessTokenMixin
     include Models::Scopes
     include Models::ResourceOwnerable
     include Models::ExpirationTimeSqlMath
+    include Models::Concerns::WriteToPrimary
 
     module ClassMethods
       # Returns an instance of the Doorkeeper::AccessToken with
@@ -66,12 +67,14 @@ def by_previous_refresh_token(previous_refresh_token)
       #   instance of the Resource Owner model or it's ID
       #
       def revoke_all_for(application_id, resource_owner, clock = Time)
-        by_resource_owner(resource_owner)
-          .where(
-            application_id: application_id,
-            revoked_at: nil,
-          )
-          .update_all(revoked_at: clock.now.utc)
+        with_primary_role do
+          by_resource_owner(resource_owner)
+            .where(
+              application_id: application_id,
+              revoked_at: nil,
+            )
+            .update_all(revoked_at: clock.now.utc)
+        end
       end
 
       # Looking for not revoked Access Token with a matching set of scopes
@@ -260,7 +263,9 @@ def create_for(application:, resource_owner:, scopes:, **token_attributes)
           token_attributes[:resource_owner_id] = resource_owner_id_for(resource_owner)
         end
 
-        create!(token_attributes)
+        with_primary_role do
+          create!(token_attributes)
+        end
       end
 
       # Looking for not revoked Access Token records that belongs to specific
@@ -435,7 +440,12 @@ def revoke_previous_refresh_token!
       return if !self.class.refresh_token_revoked_on_use? || previous_refresh_token.blank?
 
       old_refresh_token&.revoke
-      update_attribute(:previous_refresh_token, "")
+
+      if self.class.respond_to?(:with_primary_role)
+        self.class.with_primary_role { update_attribute(:previous_refresh_token, "") }
+      else
+        update_attribute(:previous_refresh_token, "")
+      end
     end
 
     private
diff --git a/lib/doorkeeper/models/concerns/revocable.rb b/lib/doorkeeper/models/concerns/revocable.rb
@@ -10,7 +10,13 @@ module Revocable
       #
       def revoke(clock = Time)
         return if revoked?
-        update_attribute(:revoked_at, clock.now.utc)
+
+        # Wrap in with_primary_role if the model class supports it
+        if self.class.respond_to?(:with_primary_role)
+          self.class.with_primary_role { update_attribute(:revoked_at, clock.now.utc) }
+        else
+          update_attribute(:revoked_at, clock.now.utc)
+        end
       end
 
       # Indicates whether the object has been revoked.
diff --git a/lib/doorkeeper/models/concerns/write_to_primary.rb b/lib/doorkeeper/models/concerns/write_to_primary.rb
@@ -0,0 +1,57 @@
+# frozen_string_literal: true
+
+module Doorkeeper
+  module Models
+    module Concerns
+      # Provides support for Rails read replicas by ensuring write operations
+      # use the primary database when automatic role switching is enabled.
+      #
+      # When Rails uses automatic role switching with read replicas, GET requests
+      # are routed to read-only databases. However, Doorkeeper may need to write
+      # to the database during GET requests (e.g., creating access tokens during
+      # implicit grant flow). This concern wraps write operations with
+      # `connected_to(role: :writing)` to ensure they use the primary database.
+      #
+      # This concern is only active when:
+      # 1. ActiveRecord supports `connected_to` (Rails 6.1+)
+      # 2. The configuration option is enabled
+      #
+      module WriteToPrimary
+        extend ActiveSupport::Concern
+
+        class_methods do
+          # Executes the given block with a connection to the primary database
+          # for writing, if read replica support is enabled and available.
+          #
+          # @yield Block to execute with write connection
+          # @return The result of the block
+          #
+          def with_primary_role(&block)
+            if should_use_primary_role?
+              ::ActiveRecord::Base.connected_to(role: :writing, &block)
+            else
+              yield
+            end
+          end
+
+          private
+
+          # Determines if we should explicitly use the primary role for writes
+          #
+          # @return [Boolean]
+          #
+          def should_use_primary_role?
+            # Guard clause: return false if ActiveRecord is not available
+            return false unless defined?(::ActiveRecord::Base)
+
+            # Only use primary role if:
+            # 1. The enable_multiple_databases option is enabled in config
+            # 2. ActiveRecord supports connected_to (Rails 6.1+)
+            Doorkeeper.config.enable_multiple_databases &&
+              ::ActiveRecord::Base.respond_to?(:connected_to)
+          end
+        end
+      end
+    end
+  end
+end
diff --git a/lib/doorkeeper/oauth/authorization/code.rb b/lib/doorkeeper/oauth/authorization/code.rb
@@ -14,7 +14,9 @@ def initialize(pre_auth, resource_owner)
         def issue_token!
           return @token if defined?(@token)
 
-          @token = Doorkeeper.config.access_grant_model.create!(access_grant_attributes)
+          @token = Doorkeeper.config.access_grant_model.with_primary_role do
+            Doorkeeper.config.access_grant_model.create!(access_grant_attributes)
+          end
         end
 
         def oob_redirect
diff --git a/lib/generators/doorkeeper/templates/initializer.rb b/lib/generators/doorkeeper/templates/initializer.rb
@@ -5,6 +5,21 @@
   # Check the list of supported ORMs here: https://github.com/doorkeeper-gem/doorkeeper#orms
   orm :active_record
 
+  # Enable support for multiple database configurations with read replicas.
+  # When enabled, Doorkeeper will wrap database write operations to ensure they
+  # use the primary (writable) database when automatic role switching is enabled.
+  #
+  # For ActiveRecord (Rails 6.1+), this uses `ActiveRecord::Base.connected_to(role: :writing)`.
+  # Other ORM extensions can implement their own primary database targeting logic.
+  #
+  # enable_multiple_databases
+  #
+  # This prevents `ActiveRecord::ReadOnlyError` when using read replicas with Rails
+  # automatic role switching. Enable this if your application uses multiple databases
+  # with automatic role switching for read replicas.
+  #
+  # See: https://guides.rubyonrails.org/active_record_multiple_databases.html#activating-automatic-role-switching
+
   # This block will be called to check whether the resource owner is authenticated or not.
   resource_owner_authenticator do
     raise "Please configure doorkeeper resource_owner_authenticator block located in #{__FILE__}"
diff --git a/spec/lib/models/concerns/write_to_primary_spec.rb b/spec/lib/models/concerns/write_to_primary_spec.rb
@@ -0,0 +1,107 @@
+# frozen_string_literal: true
+
+require "spec_helper"
+
+RSpec.describe Doorkeeper::Models::Concerns::WriteToPrimary do
+  let(:test_class) do
+    Class.new do
+      include Doorkeeper::Models::Concerns::WriteToPrimary
+
+      def self.create_record
+        with_primary_role do
+          "created"
+        end
+      end
+    end
+  end
+
+  describe ".with_primary_role" do
+    context "when ActiveRecord is not defined" do
+      before(:each) do
+        # Override the global before hook by skipping DatabaseCleaner for this context
+        # This is needed because removing ActiveRecord causes DatabaseCleaner to fail
+      end
+
+      after(:each) do
+        # Override the global after hook - don't try to clean database
+      end
+
+      around do |example|
+        # Save original ActiveRecord
+        original_active_record = Object.const_get("ActiveRecord")
+        
+        begin
+          # Temporarily hide ActiveRecord constant
+          Object.send(:remove_const, "ActiveRecord")
+          
+          # Run the test
+          Doorkeeper.configure do
+            orm :active_record
+            enable_multiple_databases
+          end
+          
+          example.run
+        ensure
+          # Restore ActiveRecord for cleanup
+          Object.const_set("ActiveRecord", original_active_record)
+        end
+      end
+
+      it "executes block without connected_to when ActiveRecord is not available" do
+        expect(test_class.create_record).to eq("created")
+      end
+    end
+
+    context "when enable_multiple_databases is disabled" do
+      before do
+        Doorkeeper.configure do
+          orm :active_record
+          # enable_multiple_databases is disabled by default
+        end
+      end
+
+      it "executes block without connected_to" do
+        expect(ActiveRecord::Base).not_to receive(:connected_to)
+        expect(test_class.create_record).to eq("created")
+      end
+    end
+
+    context "when enable_multiple_databases is enabled" do
+      before do
+        Doorkeeper.configure do
+          orm :active_record
+          enable_multiple_databases
+        end
+      end
+
+      context "when ActiveRecord supports connected_to" do
+        before do
+          allow(ActiveRecord::Base).to receive(:respond_to?)
+            .with(:connected_to)
+            .and_return(true)
+        end
+
+        it "wraps block in connected_to with writing role" do
+          expect(ActiveRecord::Base).to receive(:connected_to)
+            .with(role: :writing)
+            .and_yield
+
+          expect(test_class.create_record).to eq("created")
+        end
+      end
+
+      context "when ActiveRecord does not support connected_to" do
+        before do
+          allow(ActiveRecord::Base).to receive(:respond_to?)
+            .with(:connected_to)
+            .and_return(false)
+        end
+
+        it "executes block without connected_to" do
+          expect(ActiveRecord::Base).not_to receive(:connected_to)
+          expect(test_class.create_record).to eq("created")
+        end
+      end
+    end
+  end
+end
diff --git a/spec/lib/oauth/authorization/code_spec.rb b/spec/lib/oauth/authorization/code_spec.rb
diff --git a/spec/models/doorkeeper/access_grant_spec.rb b/spec/models/doorkeeper/access_grant_spec.rb
diff --git a/spec/models/doorkeeper/access_token_spec.rb b/spec/models/doorkeeper/access_token_spec.rb
diff --git a/spec/spec_helper.rb b/spec/spec_helper.rb
