activerecord 6.0.0 → 7.2.3

data/lib/active_record/connection_adapters/pool_manager.rb

@@ -0,0 +1,57 @@
+# frozen_string_literal: true
+
+module ActiveRecord
+  module ConnectionAdapters
+    class PoolManager # :nodoc:
+      def initialize
+        @role_to_shard_mapping = Hash.new { |h, k| h[k] = {} }
+      end
+
+      def shard_names
+        @role_to_shard_mapping.values.flat_map { |shard_map| shard_map.keys }.uniq
+      end
+
+      def role_names
+        @role_to_shard_mapping.keys
+      end
+
+      def pool_configs(role = nil)
+        if role
+          @role_to_shard_mapping[role].values
+        else
+          @role_to_shard_mapping.flat_map { |_, shard_map| shard_map.values }
+        end
+      end
+
+      def each_pool_config(role = nil, &block)
+        if role
+          @role_to_shard_mapping[role].each_value(&block)
+        else
+          @role_to_shard_mapping.each_value do |shard_map|
+            shard_map.each_value(&block)
+          end
+        end
+      end
+
+      def remove_role(role)
+        @role_to_shard_mapping.delete(role)
+      end
+
+      def remove_pool_config(role, shard)
+        @role_to_shard_mapping[role].delete(shard)
+      end
+
+      def get_pool_config(role, shard)
+        @role_to_shard_mapping[role][shard]
+      end
+
+      def set_pool_config(role, shard, pool_config)
+        if pool_config
+          @role_to_shard_mapping[role][shard] = pool_config
+        else
+          raise ArgumentError, "The `pool_config` for the :#{role} role and :#{shard} shard was `nil`. Please check your configuration. If you want your writing role to be something other than `:writing` set `config.active_record.writing_role` in your application configuration. The same setting should be applied for the `reading_role` if applicable."
+        end
+      end
+    end
+  end
+end

data/lib/active_record/connection_adapters/postgresql/column.rb

@@ -6,22 +6,74 @@ module ActiveRecord
       class Column < ConnectionAdapters::Column # :nodoc:
         delegate :oid, :fmod, to: :sql_type_metadata
 
-        def initialize(*, serial: nil, **)
+        def initialize(*, serial: nil, identity: nil, generated: nil, **)
           super
           @serial = serial
+          @identity = identity
+          @generated = generated
+        end
+
+        def identity?
+          @identity
         end
 
         def serial?
           @serial
         end
 
+        def auto_incremented_by_db?
+          serial? || identity?
+        end
+
+        def virtual?
+          # We assume every generated column is virtual, no matter the concrete type
+          @generated.present?
+        end
+
+        def has_default?
+          super && !virtual?
+        end
+
         def array
           sql_type_metadata.sql_type.end_with?("[]")
         end
         alias :array? :array
 
+        def enum?
+          type == :enum
+        end
+
         def sql_type
-          super.sub(/\[\]\z/, "")
+          super.delete_suffix("[]")
+        end
+
+        def init_with(coder)
+          @serial = coder["serial"]
+          @identity = coder["identity"]
+          @generated = coder["generated"]
+          super
+        end
+
+        def encode_with(coder)
+          coder["serial"] = @serial
+          coder["identity"] = @identity
+          coder["generated"] = @generated
+          super
+        end
+
+        def ==(other)
+          other.is_a?(Column) &&
+            super &&
+            identity? == other.identity? &&
+            serial? == other.serial?
+        end
+        alias :eql? :==
+
+        def hash
+          Column.hash ^
+            super.hash ^
+            identity?.hash ^
+            serial?.hash
         end
       end
     end

data/lib/active_record/connection_adapters/postgresql/database_statements.rb

@@ -4,132 +4,87 @@ module ActiveRecord
   module ConnectionAdapters
     module PostgreSQL
       module DatabaseStatements
-        def explain(arel, binds = [])
-          sql = "EXPLAIN #{to_sql(arel, binds)}"
-          PostgreSQL::ExplainPrettyPrinter.new.pp(exec_query(sql, "EXPLAIN", binds))
-        end
-
-        # The internal PostgreSQL identifier of the money data type.
-        MONEY_COLUMN_TYPE_OID = 790 #:nodoc:
-        # The internal PostgreSQL identifier of the BYTEA data type.
-        BYTEA_COLUMN_TYPE_OID = 17 #:nodoc:
-
-        # create a 2D array representing the result set
-        def result_as_array(res) #:nodoc:
-          # check if we have any binary column and if they need escaping
-          ftypes = Array.new(res.nfields) do |i|
-            [i, res.ftype(i)]
-          end
-
-          rows = res.values
-          return rows unless ftypes.any? { |_, x|
-            x == BYTEA_COLUMN_TYPE_OID || x == MONEY_COLUMN_TYPE_OID
-          }
-
-          typehash = ftypes.group_by { |_, type| type }
-          binaries = typehash[BYTEA_COLUMN_TYPE_OID] || []
-          monies   = typehash[MONEY_COLUMN_TYPE_OID] || []
-
-          rows.each do |row|
-            # unescape string passed BYTEA field (OID == 17)
-            binaries.each do |index, _|
-              row[index] = unescape_bytea(row[index])
-            end
-
-            # If this is a money type column and there are any currency symbols,
-            # then strip them off. Indeed it would be prettier to do this in
-            # PostgreSQLColumn.string_to_decimal but would break form input
-            # fields that call value_before_type_cast.
-            monies.each do |index, _|
-              data = row[index]
-              # Because money output is formatted according to the locale, there are two
-              # cases to consider (note the decimal separators):
-              #  (1) $12,345,678.12
-              #  (2) $12.345.678,12
-              case data
-              when /^-?\D+[\d,]+\.\d{2}$/  # (1)
-                data.gsub!(/[^-\d.]/, "")
-              when /^-?\D+[\d.]+,\d{2}$/  # (2)
-                data.gsub!(/[^-\d,]/, "").sub!(/,/, ".")
-              end
-            end
-          end
+        def explain(arel, binds = [], options = [])
+          sql    = build_explain_clause(options) + " " + to_sql(arel, binds)
+          result = internal_exec_query(sql, "EXPLAIN", binds)
+          PostgreSQL::ExplainPrettyPrinter.new.pp(result)
         end
 
         # Queries the database and returns the results in an Array-like object
-        def query(sql, name = nil) #:nodoc:
-          materialize_transactions
-
-          log(sql, name) do
-            ActiveSupport::Dependencies.interlock.permit_concurrent_loads do
-              result_as_array @connection.async_exec(sql)
+        def query(sql, name = nil) # :nodoc:
+          mark_transaction_written_if_write(sql)
+
+          log(sql, name) do |notification_payload|
+            with_raw_connection do |conn|
+              result = conn.async_exec(sql).map_types!(@type_map_for_results).values
+              verified!
+              notification_payload[:row_count] = result.count
+              result
             end
           end
         end
 
-        READ_QUERY = ActiveRecord::ConnectionAdapters::AbstractAdapter.build_read_query_regexp(:begin, :commit, :explain, :select, :set, :show, :release, :savepoint, :rollback) # :nodoc:
+        READ_QUERY = ActiveRecord::ConnectionAdapters::AbstractAdapter.build_read_query_regexp(
+          :close, :declare, :fetch, :move, :set, :show
+        ) # :nodoc:
         private_constant :READ_QUERY
 
         def write_query?(sql) # :nodoc:
           !READ_QUERY.match?(sql)
+        rescue ArgumentError # Invalid encoding
+          !READ_QUERY.match?(sql.b)
         end
 
         # Executes an SQL statement, returning a PG::Result object on success
         # or raising a PG::Error exception otherwise.
+        #
+        # Setting +allow_retry+ to true causes the db to reconnect and retry
+        # executing the SQL statement in case of a connection-related exception.
+        # This option should only be enabled for known idempotent queries.
+        #
         # 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::ReadOnlyError, "Write query attempted while in readonly mode: #{sql}"
-          end
-
-          materialize_transactions
+        def execute(...) # :nodoc:
+          super
+        ensure
+          @notice_receiver_sql_warnings = []
+        end
 
-          log(sql, name) do
-            ActiveSupport::Dependencies.interlock.permit_concurrent_loads do
-              @connection.async_exec(sql)
+        def raw_execute(sql, name, async: false, allow_retry: false, materialize_transactions: true)
+          log(sql, name, async: async) do |notification_payload|
+            with_raw_connection(allow_retry: allow_retry, materialize_transactions: materialize_transactions) do |conn|
+              result = conn.async_exec(sql)
+              verified!
+              handle_warnings(result)
+              notification_payload[:row_count] = result.count
+              result
             end
           end
         end
 
-        def exec_query(sql, name = "SQL", binds = [], prepare: false)
-          execute_and_clear(sql, name, binds, prepare: prepare) do |result|
+        def internal_exec_query(sql, name = "SQL", binds = [], prepare: false, async: false, allow_retry: false, materialize_transactions: true) # :nodoc:
+          execute_and_clear(sql, name, binds, prepare: prepare, async: async, allow_retry: allow_retry, materialize_transactions: materialize_transactions) do |result|
             types = {}
             fields = result.fields
             fields.each_with_index do |fname, i|
               ftype = result.ftype i
               fmod  = result.fmod i
-              types[fname] = get_oid_type(ftype, fmod, fname)
+              types[fname] = types[i] = get_oid_type(ftype, fmod, fname)
             end
-            ActiveRecord::Result.new(fields, result.values, types)
+            build_result(columns: fields, rows: result.values, column_types: types.freeze)
           end
         end
 
-        def exec_delete(sql, name = nil, binds = [])
+        def exec_delete(sql, name = nil, binds = []) # :nodoc:
           execute_and_clear(sql, name, binds) { |result| result.cmd_tuples }
         end
         alias :exec_update :exec_delete
 
-        def sql_for_insert(sql, pk, binds) # :nodoc:
-          if pk.nil?
-            # Extract the table from the insert sql. Yuck.
-            table_ref = extract_table_ref_from_insert_sql(sql)
-            pk = primary_key(table_ref) if table_ref
-          end
-
-          if pk = suppress_composite_primary_key(pk)
-            sql = "#{sql} RETURNING #{quote_column_name(pk)}"
-          end
-
-          super
-        end
-        private :sql_for_insert
-
-        def exec_insert(sql, name = nil, binds = [], pk = nil, sequence_name = nil)
+        def exec_insert(sql, name = nil, binds = [], pk = nil, sequence_name = nil, returning: nil) # :nodoc:
           if use_insert_returning? || pk == false
             super
           else
-            result = exec_query(sql, name, binds)
+            result = internal_exec_query(sql, name, binds)
             unless sequence_name
               table_ref = extract_table_ref_from_insert_sql(sql)
               if table_ref
@@ -144,38 +99,110 @@ module ActiveRecord
         end
 
         # Begins a transaction.
-        def begin_db_transaction
-          execute "BEGIN"
+        def begin_db_transaction # :nodoc:
+          internal_execute("BEGIN", "TRANSACTION", allow_retry: true, materialize_transactions: false)
         end
 
-        def begin_isolated_db_transaction(isolation)
-          begin_db_transaction
-          execute "SET TRANSACTION ISOLATION LEVEL #{transaction_isolation_levels.fetch(isolation)}"
+        def begin_isolated_db_transaction(isolation) # :nodoc:
+          internal_execute("BEGIN ISOLATION LEVEL #{transaction_isolation_levels.fetch(isolation)}", "TRANSACTION", allow_retry: true, materialize_transactions: false)
         end
 
         # Commits a transaction.
-        def commit_db_transaction
-          execute "COMMIT"
+        def commit_db_transaction # :nodoc:
+          internal_execute("COMMIT", "TRANSACTION", allow_retry: false, materialize_transactions: true)
         end
 
         # Aborts a transaction.
-        def exec_rollback_db_transaction
-          execute "ROLLBACK"
+        def exec_rollback_db_transaction # :nodoc:
+          cancel_any_running_query
+          internal_execute("ROLLBACK", "TRANSACTION", allow_retry: false, materialize_transactions: true)
+        end
+
+        def exec_restart_db_transaction # :nodoc:
+          cancel_any_running_query
+          internal_execute("ROLLBACK AND CHAIN", "TRANSACTION", allow_retry: false, materialize_transactions: true)
+        end
+
+        # From https://www.postgresql.org/docs/current/functions-datetime.html#FUNCTIONS-DATETIME-CURRENT
+        HIGH_PRECISION_CURRENT_TIMESTAMP = Arel.sql("CURRENT_TIMESTAMP", retryable: true).freeze # :nodoc:
+        private_constant :HIGH_PRECISION_CURRENT_TIMESTAMP
+
+        def high_precision_current_timestamp
+          HIGH_PRECISION_CURRENT_TIMESTAMP
+        end
+
+        def build_explain_clause(options = [])
+          return "EXPLAIN" if options.empty?
+
+          "EXPLAIN (#{options.join(", ").upcase})"
+        end
+
+        # Set when constraints will be checked for the current transaction.
+        #
+        # Not passing any specific constraint names will set the value for all deferrable constraints.
+        #
+        # [<tt>deferred</tt>]
+        #   Valid values are +:deferred+ or +:immediate+.
+        #
+        # See https://www.postgresql.org/docs/current/sql-set-constraints.html
+        def set_constraints(deferred, *constraints)
+          unless %i[deferred immediate].include?(deferred)
+            raise ArgumentError, "deferred must be :deferred or :immediate"
+          end
+
+          constraints = if constraints.empty?
+            "ALL"
+          else
+            constraints.map { |c| quote_table_name(c) }.join(", ")
+          end
+          execute("SET CONSTRAINTS #{constraints} #{deferred.to_s.upcase}")
         end
 
         private
-          def build_truncate_statements(*table_names)
-            "TRUNCATE TABLE #{table_names.map(&method(:quote_table_name)).join(", ")}"
+          IDLE_TRANSACTION_STATUSES = [PG::PQTRANS_IDLE, PG::PQTRANS_INTRANS, PG::PQTRANS_INERROR]
+          private_constant :IDLE_TRANSACTION_STATUSES
+
+          def cancel_any_running_query
+            return if @raw_connection.nil? || IDLE_TRANSACTION_STATUSES.include?(@raw_connection.transaction_status)
+
+            @raw_connection.cancel
+            @raw_connection.block
+          rescue PG::Error
+          end
+
+          def execute_batch(statements, name = nil)
+            execute(combine_multi_statements(statements))
+          end
+
+          def build_truncate_statements(table_names)
+            ["TRUNCATE TABLE #{table_names.map(&method(:quote_table_name)).join(", ")}"]
           end
 
           # Returns the current ID of a table's sequence.
           def last_insert_id_result(sequence_name)
-            exec_query("SELECT currval(#{quote(sequence_name)})", "SQL")
+            internal_exec_query("SELECT currval(#{quote(sequence_name)})", "SQL")
+          end
+
+          def returning_column_values(result)
+            result.rows.first
           end
 
           def suppress_composite_primary_key(pk)
             pk unless pk.is_a?(Array)
           end
+
+          def handle_warnings(sql)
+            @notice_receiver_sql_warnings.each do |warning|
+              next if warning_ignored?(warning)
+
+              warning.sql = sql
+              ActiveRecord.db_warnings_action.call(warning)
+            end
+          end
+
+          def warning_ignored?(warning)
+            ["WARNING", "ERROR", "FATAL", "PANIC"].exclude?(warning.level) || super
+          end
       end
     end
   end

data/lib/active_record/connection_adapters/postgresql/oid/array.rb

@@ -65,7 +65,7 @@ module ActiveRecord
           end
 
           def map(value, &block)
-            value.map(&block)
+            value.map { |v| subtype.map(v, &block) }
           end
 
           def changed_in_place?(raw_old_value, new_value)
@@ -77,7 +77,6 @@ module ActiveRecord
           end
 
           private
-
             def type_cast_array(value, method)
               if value.is_a?(::Array)
                 value.map { |item| type_cast_array(item, method) }

data/lib/active_record/connection_adapters/postgresql/oid/cidr.rb

@@ -12,24 +12,28 @@ module ActiveRecord
           end
 
           def type_cast_for_schema(value)
-            subnet_mask = value.instance_variable_get(:@mask_addr)
-
             # If the subnet mask is equal to /32, don't output it
-            if subnet_mask == (2**32 - 1)
+            if value.prefix == 32
               "\"#{value}\""
             else
-              "\"#{value}/#{subnet_mask.to_s(2).count('1')}\""
+              "\"#{value}/#{value.prefix}\""
             end
           end
 
           def serialize(value)
             if IPAddr === value
-              "#{value}/#{value.instance_variable_get(:@mask_addr).to_s(2).count('1')}"
+              "#{value}/#{value.prefix}"
             else
               value
             end
           end
 
+          # TODO: Remove when IPAddr#== compares IPAddr#prefix. See
+          # https://github.com/ruby/ipaddr/issues/21
+          def changed?(old_value, new_value, _new_value_before_type_cast)
+            !old_value.eql?(new_value) || !old_value.nil? && old_value.prefix != new_value.prefix
+          end
+
           def cast_value(value)
             if value.nil?
               nil

data/lib/active_record/connection_adapters/postgresql/oid/date.rb

@@ -10,12 +10,20 @@ module ActiveRecord
             when "infinity" then ::Float::INFINITY
             when "-infinity" then -::Float::INFINITY
             when / BC$/
-              astronomical_year = format("%04d", -value[/^\d+/].to_i + 1)
-              super(value.sub(/ BC$/, "").sub(/^\d+/, astronomical_year))
+              value = value.sub(/^\d+/) { |year| format("%04d", -year.to_i + 1) }
+              super(value.delete_suffix!(" BC"))
             else
               super
             end
           end
+
+          def type_cast_for_schema(value)
+            case value
+            when ::Float::INFINITY then "::Float::INFINITY"
+            when -::Float::INFINITY then "-::Float::INFINITY"
+            else super
+            end
+          end
         end
       end
     end

data/lib/active_record/connection_adapters/postgresql/oid/date_time.rb

@@ -10,12 +10,25 @@ module ActiveRecord
             when "infinity" then ::Float::INFINITY
             when "-infinity" then -::Float::INFINITY
             when / BC$/
-              astronomical_year = format("%04d", -value[/^\d+/].to_i + 1)
-              super(value.sub(/ BC$/, "").sub(/^\d+/, astronomical_year))
+              value = value.sub(/^\d+/) { |year| format("%04d", -year.to_i + 1) }
+              super(value.delete_suffix!(" BC"))
             else
               super
             end
           end
+
+          def type_cast_for_schema(value)
+            case value
+            when ::Float::INFINITY then "::Float::INFINITY"
+            when -::Float::INFINITY then "-::Float::INFINITY"
+            else super
+            end
+          end
+
+          protected
+            def real_type_unless_aliased(real_type)
+              ActiveRecord::ConnectionAdapters::PostgreSQLAdapter.datetime_type == real_type ? :datetime : real_type
+            end
         end
       end
     end

data/lib/active_record/connection_adapters/postgresql/oid/enum.rb

@@ -10,7 +10,6 @@ module ActiveRecord
           end
 
           private
-
             def cast_value(value)
               value.to_s
             end

data/lib/active_record/connection_adapters/postgresql/oid/hstore.rb

@@ -1,10 +1,14 @@
 # frozen_string_literal: true
 
+require "strscan"
+
 module ActiveRecord
   module ConnectionAdapters
     module PostgreSQL
       module OID # :nodoc:
         class Hstore < Type::Value # :nodoc:
+          ERROR = "Invalid Hstore document: %s"
+
           include ActiveModel::Type::Helpers::Mutable
 
           def type
@@ -12,15 +16,56 @@ module ActiveRecord
           end
 
           def deserialize(value)
-            if value.is_a?(::String)
-              ::Hash[value.scan(HstorePair).map { |k, v|
-                v = v.upcase == "NULL" ? nil : v.gsub(/\A"(.*)"\Z/m, '\1').gsub(/\\(.)/, '\1')
-                k = k.gsub(/\A"(.*)"\Z/m, '\1').gsub(/\\(.)/, '\1')
-                [k, v]
-              }]
-            else
-              value
+            return value unless value.is_a?(::String)
+
+            scanner = StringScanner.new(value)
+            hash = {}
+
+            until scanner.eos?
+              unless scanner.skip(/"/)
+                raise(ArgumentError, ERROR % scanner.string.inspect)
+              end
+
+              unless key = scanner.scan(/^(\\[\\"]|[^\\"])*?(?=")/)
+                raise(ArgumentError, ERROR % scanner.string.inspect)
+              end
+
+              unless scanner.skip(/"=>?/)
+                raise(ArgumentError, ERROR % scanner.string.inspect)
+              end
+
+              if scanner.scan(/NULL/)
+                value = nil
+              else
+                unless scanner.skip(/"/)
+                  raise(ArgumentError, ERROR % scanner.string.inspect)
+                end
+
+                unless value = scanner.scan(/^(\\[\\"]|[^\\"])*?(?=")/)
+                  raise(ArgumentError, ERROR % scanner.string.inspect)
+                end
+
+                unless scanner.skip(/"/)
+                  raise(ArgumentError, ERROR % scanner.string.inspect)
+                end
+              end
+
+              key.gsub!('\"', '"')
+              key.gsub!("\\\\", "\\")
+
+              if value
+                value.gsub!('\"', '"')
+                value.gsub!("\\\\", "\\")
+              end
+
+              hash[key] = value
+
+              unless scanner.skip(/, /) || scanner.eos?
+                raise(ArgumentError, ERROR % scanner.string.inspect)
+              end
             end
+
+            hash
           end
 
           def serialize(value)
@@ -46,13 +91,6 @@ module ActiveRecord
           end
 
           private
-
-            HstorePair = begin
-              quoted_string = /"[^"\\]*(?:\\.[^"\\]*)*"/
-              unquoted_string = /(?:\\.|[^\s,])[^\s=,\\]*(?:\\.[^\s=,\\]*|=[^,>])*/
-              /(#{quoted_string}|#{unquoted_string})\s*=>\s*(#{quoted_string}|#{unquoted_string})/
-            end
-
             def escape_hstore(value)
               if value.nil?
                 "NULL"