Merge pull request #34505 from eileencodes/add-readonly-mode

Add ability to block writes to a database

5 files changed, 64 insertions, 1 deletions

diff --git a/activerecord/lib/active_record/connection_adapters/abstract/database_statements.rb b/activerecord/lib/active_record/connection_adapters/abstract/database_statements.rb
index 0059f0b773..2299fc0214 100644
--- a/activerecord/lib/active_record/connection_adapters/abstract/database_statements.rb
+++ b/activerecord/lib/active_record/connection_adapters/abstract/database_statements.rb
@@ -98,6 +98,11 @@ module ActiveRecord
         exec_query(sql, name).rows
       end
 
+      # Determines whether the SQL statement is a write query.
+      def write_query?(sql)
+        raise NotImplementedError
+      end
+
       # Executes the SQL statement in the context of this connection and returns
       # the raw result from the connection adapter.
       # Note: depending on your database connector, the result returned by this
diff --git a/activerecord/lib/active_record/connection_adapters/abstract_adapter.rb b/activerecord/lib/active_record/connection_adapters/abstract_adapter.rb
index f73369ceef..31a994292e 100644
--- a/activerecord/lib/active_record/connection_adapters/abstract_adapter.rb
+++ b/activerecord/lib/active_record/connection_adapters/abstract_adapter.rb
@@ -76,7 +76,7 @@ module ActiveRecord
 
       SIMPLE_INT = /\A\d+\z/
 
-      attr_accessor :visitor, :pool
+      attr_accessor :visitor, :pool, :prevent_writes
       attr_reader :schema_cache, :owner, :logger, :prepared_statements, :lock
       alias :in_use? :owner
 
@@ -100,6 +100,13 @@ module ActiveRecord
         end
       end
 
+      def self.build_read_query_regexp(*parts) # :nodoc:
+        lambda do |*parts|
+          parts = parts.map { |part| /\A\s*#{part}/i }
+          Regexp.union(*parts)
+        end
+      end
+
       def initialize(connection, logger = nil, config = {}) # :nodoc:
         super()
 
@@ -133,6 +140,27 @@ module ActiveRecord
         @config[:replica] || false
       end
 
+      # Determines whether writes are currently being prevents.
+      #
+      # Returns true if the connection is a replica, or if +prevent_writes+
+      # is set to true.
+      def preventing_writes?
+        replica? || prevent_writes
+      end
+
+      # Prevent writing to the database  regardless of role.
+      #
+      # In some cases you may want to prevent writes to the database
+      # even if you are on a database that can write. `while_preventing_writes`
+      # will prevent writes to the database for the duration of the block.
+      def while_preventing_writes
+        original = self.prevent_writes
+        self.prevent_writes = true
+        yield
+      ensure
+        self.prevent_writes = original
+      end
+
       def migrations_paths # :nodoc:
         @config[:migrations_paths] || Migrator.migrations_paths
       end
diff --git a/activerecord/lib/active_record/connection_adapters/mysql/database_statements.rb b/activerecord/lib/active_record/connection_adapters/mysql/database_statements.rb
index 9ea237dd81..0ec9581c42 100644
--- a/activerecord/lib/active_record/connection_adapters/mysql/database_statements.rb
+++ b/activerecord/lib/active_record/connection_adapters/mysql/database_statements.rb
@@ -19,8 +19,18 @@ module ActiveRecord
           execute(sql, name).to_a
         end
 
+        READ_QUERY = ActiveRecord::ConnectionAdapters::AbstractAdapter.build_read_query_regexp.call(:begin, :select, :set, :show, :release, :savepoint) # :nodoc:
+
+        def write_query?(sql) # :nodoc:
+          !READ_QUERY.match?(sql)
+        end
+
         # Executes the SQL statement in the context of this connection.
         def execute(sql, name = nil)
+          if preventing_writes? && write_query?(sql)
+            raise ActiveRecord::StatementInvalid, "Write query attempted while in readonly mode: #{sql}"
+          end
+
           # make sure we carry over any changes to ActiveRecord::Base.default_timezone that have been
           # made since we established the connection
           @connection.query_options[:database_timezone] = ActiveRecord::Base.default_timezone
diff --git a/activerecord/lib/active_record/connection_adapters/postgresql/database_statements.rb b/activerecord/lib/active_record/connection_adapters/postgresql/database_statements.rb
index 6bd6b67165..c8bc339e61 100644
--- a/activerecord/lib/active_record/connection_adapters/postgresql/database_statements.rb
+++ b/activerecord/lib/active_record/connection_adapters/postgresql/database_statements.rb
@@ -67,11 +67,21 @@ module ActiveRecord
           end
         end
 
+        READ_QUERY = ActiveRecord::ConnectionAdapters::AbstractAdapter.build_read_query_regexp.call(:select, :show, :set) # :nodoc:
+
+        def write_query?(sql) # :nodoc:
+          !READ_QUERY.match?(sql)
+        end
+
         # Executes an SQL statement, returning a PG::Result object on success
         # or raising a PG::Error exception otherwise.
         # Note: the PG::Result object is manually memory managed; if you don't
         # need it specifically, you may want consider the <tt>exec_query</tt> wrapper.
         def execute(sql, name = nil)
+          if preventing_writes? && write_query?(sql)
+            raise ActiveRecord::StatementInvalid, "Write query attempted while in readonly mode: #{sql}"
+          end
+
           materialize_transactions
 
           log(sql, name) do
diff --git a/activerecord/lib/active_record/connection_adapters/sqlite3_adapter.rb b/activerecord/lib/active_record/connection_adapters/sqlite3_adapter.rb
index b41bf2fc66..1bec8fbabd 100644
--- a/activerecord/lib/active_record/connection_adapters/sqlite3_adapter.rb
+++ b/activerecord/lib/active_record/connection_adapters/sqlite3_adapter.rb
@@ -209,6 +209,12 @@ module ActiveRecord
       # DATABASE STATEMENTS ======================================
       #++
 
+      READ_QUERY = ActiveRecord::ConnectionAdapters::AbstractAdapter.build_read_query_regexp.call(:select) # :nodoc:
+
+      def write_query?(sql) # :nodoc:
+        !READ_QUERY.match?(sql)
+      end
+
       def explain(arel, binds = [])
         sql = "EXPLAIN QUERY PLAN #{to_sql(arel, binds)}"
         SQLite3::ExplainPrettyPrinter.new.pp(exec_query(sql, "EXPLAIN", []))
@@ -257,6 +263,10 @@ module ActiveRecord
       end
 
       def execute(sql, name = nil) #:nodoc:
+        if preventing_writes? && write_query?(sql)
+          raise ActiveRecord::StatementInvalid, "Write query attempted while in readonly mode: #{sql}"
+        end
+
         materialize_transactions
 
         log(sql, name) do