ActiveRecord adapter: wrap all reads/writes in with_connection

lib/flipper/adapters/active_record.rb

@@ -53,18 +53,20 @@ def initialize(options = {})
 
       # Public: The set of known features.
       def features
-        @feature_class.all.map(&:key).to_set
+        with_connection(@feature_class) { @feature_class.all.map(&:key).to_set }
       end
 
       # Public: Adds a feature to the set of known features.
       def add(feature)
-        # race condition, but add is only used by enable/disable which happen
-        # super rarely, so it shouldn't matter in practice
-        @feature_class.transaction do
-          unless @feature_class.where(key: feature.key).first
-            begin
-              @feature_class.create! { |f| f.key = feature.key }
-            rescue ::ActiveRecord::RecordNotUnique
+        with_connection(@feature_class) do
+          # race condition, but add is only used by enable/disable which happen
+          # super rarely, so it shouldn't matter in practice
+          @feature_class.transaction do
+            unless @feature_class.where(key: feature.key).first
+              begin
+                @feature_class.create! { |f| f.key = feature.key }
+              rescue ::ActiveRecord::RecordNotUnique
+              end
             end
           end
         end
@@ -74,52 +76,58 @@ def add(feature)
 
       # Public: Removes a feature from the set of known features.
       def remove(feature)
-        @feature_class.transaction do
-          @feature_class.where(key: feature.key).destroy_all
-          clear(feature)
+        with_connection(@feature_class) do
+          @feature_class.transaction do
+            @feature_class.where(key: feature.key).destroy_all
+            clear(feature)
+          end
         end
         true
       end
 
       # Public: Clears the gate values for a feature.
       def clear(feature)
-        @gate_class.where(feature_key: feature.key).destroy_all
+        with_connection(@gate_class) { @gate_class.where(feature_key: feature.key).destroy_all }
         true
       end
 
       # Public: Gets the values for all gates for a given feature.
       #
       # Returns a Hash of Flipper::Gate#key => value.
       def get(feature)
-        db_gates = @gate_class.where(feature_key: feature.key)
+        db_gates = with_connection(@gate_class) { @gate_class.where(feature_key: feature.key) }
         result_for_feature(feature, db_gates)
       end
 
       def get_multi(features)
-        db_gates = @gate_class.where(feature_key: features.map(&:key))
-        grouped_db_gates = db_gates.group_by(&:feature_key)
-        result = {}
-        features.each do |feature|
-          result[feature.key] = result_for_feature(feature, grouped_db_gates[feature.key])
+        with_connection(@gate_class) do
+          db_gates = @gate_class.where(feature_key: features.map(&:key))
+          grouped_db_gates = db_gates.group_by(&:feature_key)
+          result = {}
+          features.each do |feature|
+            result[feature.key] = result_for_feature(feature, grouped_db_gates[feature.key])
+          end
+          result
         end
-        result
       end
 
       def get_all
-        features = ::Arel::Table.new(@feature_class.table_name.to_sym)
-        gates = ::Arel::Table.new(@gate_class.table_name.to_sym)
-        rows_query = features.join(gates, Arel::Nodes::OuterJoin)
-          .on(features[:key].eq(gates[:feature_key]))
-          .project(features[:key].as('feature_key'), gates[:key], gates[:value])
-        rows = @feature_class.connection.select_all rows_query
-        db_gates = rows.map { |row| @gate_class.new(row) }
-        grouped_db_gates = db_gates.group_by(&:feature_key)
-        result = Hash.new { |hash, key| hash[key] = default_config }
-        features = grouped_db_gates.keys.map { |key| Flipper::Feature.new(key, self) }
-        features.each do |feature|
-          result[feature.key] = result_for_feature(feature, grouped_db_gates[feature.key])
+        with_connection(@feature_class) do
+          features = ::Arel::Table.new(@feature_class.table_name.to_sym)
+          gates = ::Arel::Table.new(@gate_class.table_name.to_sym)
+          rows_query = features.join(gates, Arel::Nodes::OuterJoin)
+            .on(features[:key].eq(gates[:feature_key]))
+            .project(features[:key].as('feature_key'), gates[:key], gates[:value])
+          rows = @feature_class.connection.select_all rows_query
+          db_gates = rows.map { |row| @gate_class.new(row) }
+          grouped_db_gates = db_gates.group_by(&:feature_key)
+          result = Hash.new { |hash, key| hash[key] = default_config }
+          features = grouped_db_gates.keys.map { |key| Flipper::Feature.new(key, self) }
+          features.each do |feature|
+            result[feature.key] = result_for_feature(feature, grouped_db_gates[feature.key])
+          end
+          result
         end
-        result
       end
 
       # Public: Enables a gate for a given thing.
@@ -158,7 +166,9 @@ def disable(feature, gate, thing)
         when :integer
           set(feature, gate, thing)
         when :set
-          @gate_class.where(feature_key: feature.key, key: gate.key, value: thing.value).destroy_all
+          with_connection(@gate_class) do
+            @gate_class.where(feature_key: feature.key, key: gate.key, value: thing.value).destroy_all
+          end
         else
           unsupported_data_type gate.data_type
         end
@@ -175,29 +185,33 @@ def unsupported_data_type(data_type)
 
       def set(feature, gate, thing, options = {})
         clear_feature = options.fetch(:clear, false)
-        @gate_class.transaction do
-          clear(feature) if clear_feature
-          @gate_class.where(feature_key: feature.key, key: gate.key).destroy_all
-          begin
-            @gate_class.create! do |g|
-              g.feature_key = feature.key
-              g.key = gate.key
-              g.value = thing.value.to_s
+        with_connection(@gate_class) do
+          @gate_class.transaction do
+            clear(feature) if clear_feature
+            @gate_class.where(feature_key: feature.key, key: gate.key).destroy_all
+            begin
+              @gate_class.create! do |g|
+                g.feature_key = feature.key
+                g.key = gate.key
+                g.value = thing.value.to_s
+              end
+            rescue ::ActiveRecord::RecordNotUnique
+              # assume this happened concurrently with the same thing and its fine
+              # see https://github.com/jnunemaker/flipper/issues/544
             end
-          rescue ::ActiveRecord::RecordNotUnique
-            # assume this happened concurrently with the same thing and its fine
-            # see https://github.com/jnunemaker/flipper/issues/544
           end
         end
 
         nil
       end
 
       def enable_multi(feature, gate, thing)
-        @gate_class.create! do |g|
-          g.feature_key = feature.key
-          g.key = gate.key
-          g.value = thing.value.to_s
+        with_connection(@gate_class) do
+          @gate_class.create! do |g|
+            g.feature_key = feature.key
+            g.key = gate.key
+            g.value = thing.value.to_s
+          end
         end
 
         nil
@@ -227,6 +241,10 @@ def result_for_feature(feature, db_gates)
         end
         result
       end
+
+      def with_connection(model = @feature_class, &block)
+        model.connection_pool.with_connection(&block)
+      end
     end
   end
 end

spec/flipper/adapters/active_record_spec.rb

@@ -57,4 +57,126 @@
       expect(Flipper.adapter.adapter).to be_a(Flipper::Adapters::ActiveRecord)
     end
   end
+
+  context "ActiveRecord connection_pool" do
+    before do
+      ActiveRecord::Base.clear_active_connections!
+    end
+
+    context "#features" do
+      it "does not hold onto connections" do
+        expect(ActiveRecord::Base.connection_handler.active_connections?).to be(false)
+        subject.features
+        expect(ActiveRecord::Base.connection_handler.active_connections?).to be(false)
+      end
+
+      it "does not release previously held connection" do
+        ActiveRecord::Base.connection # establish a new connection
+        expect(ActiveRecord::Base.connection_handler.active_connections?).to be(true)
+        subject.features
+        expect(ActiveRecord::Base.connection_handler.active_connections?).to be(true)
+      end
+    end
+
+    context "#get_all" do
+      it "does not hold onto connections" do
+        expect(ActiveRecord::Base.connection_handler.active_connections?).to be(false)
+        subject.get_all
+        expect(ActiveRecord::Base.connection_handler.active_connections?).to be(false)
+      end
+
+      it "does not release previously held connection" do
+        ActiveRecord::Base.connection # establish a new connection
+        expect(ActiveRecord::Base.connection_handler.active_connections?).to be(true)
+        subject.get_all
+        expect(ActiveRecord::Base.connection_handler.active_connections?).to be(true)
+      end
+    end
+
+    context "#add / #remove / #clear" do
+      let(:feature) { Flipper::Feature.new(:search, subject) }
+
+      it "does not hold onto connections" do
+        expect(ActiveRecord::Base.connection_handler.active_connections?).to be(false)
+        subject.add(feature)
+        expect(ActiveRecord::Base.connection_handler.active_connections?).to be(false)
+        subject.remove(feature)
+        expect(ActiveRecord::Base.connection_handler.active_connections?).to be(false)
+        subject.clear(feature)
+        expect(ActiveRecord::Base.connection_handler.active_connections?).to be(false)
+      end
+
+      it "does not release previously held connection" do
+        ActiveRecord::Base.connection # establish a new connection
+        expect(ActiveRecord::Base.connection_handler.active_connections?).to be(true)
+        subject.add(feature)
+        expect(ActiveRecord::Base.connection_handler.active_connections?).to be(true)
+        subject.remove(feature)
+        expect(ActiveRecord::Base.connection_handler.active_connections?).to be(true)
+        subject.clear(feature)
+        expect(ActiveRecord::Base.connection_handler.active_connections?).to be(true)
+      end
+    end
+
+    context "#get_multi" do
+      let(:feature) { Flipper::Feature.new(:search, subject) }
+
+      it "does not hold onto connections" do
+        expect(ActiveRecord::Base.connection_handler.active_connections?).to be(false)
+        subject.get_multi([feature])
+        expect(ActiveRecord::Base.connection_handler.active_connections?).to be(false)
+      end
+
+      it "does not release previously held connection" do
+        ActiveRecord::Base.connection # establish a new connection
+        expect(ActiveRecord::Base.connection_handler.active_connections?).to be(true)
+        subject.get_multi([feature])
+        expect(ActiveRecord::Base.connection_handler.active_connections?).to be(true)
+      end
+    end
+
+    context "#enable/#disable boolean" do
+      let(:feature) { Flipper::Feature.new(:search, subject) }
+      let(:gate) { feature.gate(:boolean)}
+
+      it "does not hold onto connections" do
+        expect(ActiveRecord::Base.connection_handler.active_connections?).to be(false)
+        subject.enable(feature, gate, gate.wrap(true))
+        expect(ActiveRecord::Base.connection_handler.active_connections?).to be(false)
+        subject.disable(feature, gate, gate.wrap(false))
+        expect(ActiveRecord::Base.connection_handler.active_connections?).to be(false)
+      end
+
+      it "does not release previously held connection" do
+        ActiveRecord::Base.connection # establish a new connection
+        expect(ActiveRecord::Base.connection_handler.active_connections?).to be(true)
+        subject.enable(feature, gate, gate.wrap(true))
+        expect(ActiveRecord::Base.connection_handler.active_connections?).to be(true)
+        subject.disable(feature, gate, gate.wrap(false))
+        expect(ActiveRecord::Base.connection_handler.active_connections?).to be(true)
+      end
+    end
+
+    context "#enable/#disable set" do
+      let(:feature) { Flipper::Feature.new(:search, subject) }
+      let(:gate) { feature.gate(:group) }
+
+      it "does not hold onto connections" do
+        expect(ActiveRecord::Base.connection_handler.active_connections?).to be(false)
+        subject.enable(feature, gate, gate.wrap(:admin))
+        expect(ActiveRecord::Base.connection_handler.active_connections?).to be(false)
+        subject.disable(feature, gate, gate.wrap(:admin))
+        expect(ActiveRecord::Base.connection_handler.active_connections?).to be(false)
+      end
+
+      it "does not release previously held connection" do
+        ActiveRecord::Base.connection # establish a new connection
+        expect(ActiveRecord::Base.connection_handler.active_connections?).to be(true)
+        subject.enable(feature, gate, gate.wrap(:admin))
+        expect(ActiveRecord::Base.connection_handler.active_connections?).to be(true)
+        subject.disable(feature, gate, gate.wrap(:admin))
+        expect(ActiveRecord::Base.connection_handler.active_connections?).to be(true)
+      end
+    end
+  end
 end