omg-activerecord 8.0.0.alpha1

data/lib/active_record/connection_adapters/mysql2_adapter.rb

@@ -0,0 +1,196 @@
+# frozen_string_literal: true
+
+require "active_record/connection_adapters/abstract_mysql_adapter"
+require "active_record/connection_adapters/mysql2/database_statements"
+
+gem "mysql2", "~> 0.5"
+require "mysql2"
+
+module ActiveRecord
+  module ConnectionAdapters
+    # = Active Record MySQL2 Adapter
+    class Mysql2Adapter < AbstractMysqlAdapter
+      ER_BAD_DB_ERROR           = 1049
+      ER_DBACCESS_DENIED_ERROR  = 1044
+      ER_ACCESS_DENIED_ERROR    = 1045
+      ER_CONN_HOST_ERROR        = 2003
+      ER_UNKNOWN_HOST_ERROR     = 2005
+
+      ADAPTER_NAME = "Mysql2"
+
+      include Mysql2::DatabaseStatements
+
+      class << self
+        def new_client(config)
+          ::Mysql2::Client.new(config)
+        rescue ::Mysql2::Error => error
+          case error.error_number
+          when ER_BAD_DB_ERROR
+            raise ActiveRecord::NoDatabaseError.db_error(config[:database])
+          when ER_DBACCESS_DENIED_ERROR, ER_ACCESS_DENIED_ERROR
+            raise ActiveRecord::DatabaseConnectionError.username_error(config[:username])
+          when ER_CONN_HOST_ERROR, ER_UNKNOWN_HOST_ERROR
+            raise ActiveRecord::DatabaseConnectionError.hostname_error(config[:host])
+          else
+            raise ActiveRecord::ConnectionNotEstablished, error.message
+          end
+        end
+
+        private
+          def initialize_type_map(m)
+            super
+
+            m.register_type(%r(char)i) do |sql_type|
+              limit = extract_limit(sql_type)
+              Type.lookup(:string, adapter: :mysql2, limit: limit)
+            end
+
+            m.register_type %r(^enum)i, Type.lookup(:string, adapter: :mysql2)
+            m.register_type %r(^set)i,  Type.lookup(:string, adapter: :mysql2)
+          end
+      end
+
+      TYPE_MAP = Type::TypeMap.new.tap { |m| initialize_type_map(m) }
+
+      def initialize(...)
+        super
+
+        @affected_rows_before_warnings = nil
+        @config[:flags] ||= 0
+
+        if @config[:flags].kind_of? Array
+          @config[:flags].push "FOUND_ROWS"
+        else
+          @config[:flags] |= ::Mysql2::Client::FOUND_ROWS
+        end
+
+        @connection_parameters ||= @config
+      end
+
+      def supports_json?
+        !mariadb? && database_version >= "5.7.8"
+      end
+
+      def supports_comments?
+        true
+      end
+
+      def supports_comments_in_create?
+        true
+      end
+
+      def supports_savepoints?
+        true
+      end
+
+      def savepoint_errors_invalidate_transactions?
+        true
+      end
+
+      def supports_lazy_transactions?
+        true
+      end
+
+      # HELPER METHODS ===========================================
+
+      def error_number(exception)
+        exception.error_number if exception.respond_to?(:error_number)
+      end
+
+      #--
+      # CONNECTION MANAGEMENT ====================================
+      #++
+
+      def connected?
+        !(@raw_connection.nil? || @raw_connection.closed?)
+      end
+
+      def active?
+        connected? && @lock.synchronize { @raw_connection&.ping } || false
+      end
+
+      alias :reset! :reconnect!
+
+      # Disconnects from the database if already connected.
+      # Otherwise, this method does nothing.
+      def disconnect!
+        @lock.synchronize do
+          super
+          @raw_connection&.close
+          @raw_connection = nil
+        end
+      end
+
+      def discard! # :nodoc:
+        @lock.synchronize do
+          super
+          @raw_connection&.automatic_close = false
+          @raw_connection = nil
+        end
+      end
+
+      private
+        def text_type?(type)
+          TYPE_MAP.lookup(type).is_a?(Type::String) || TYPE_MAP.lookup(type).is_a?(Type::Text)
+        end
+
+        def connect
+          @raw_connection = self.class.new_client(@connection_parameters)
+        rescue ConnectionNotEstablished => ex
+          raise ex.set_pool(@pool)
+        end
+
+        def reconnect
+          @lock.synchronize do
+            @raw_connection&.close
+            @raw_connection = nil
+            connect
+          end
+        end
+
+        def configure_connection
+          @raw_connection.query_options[:as] = :array
+          @raw_connection.query_options[:database_timezone] = default_timezone
+          super
+        end
+
+        def full_version
+          database_version.full_version_string
+        end
+
+        def get_full_version
+          any_raw_connection.server_info[:version]
+        end
+
+        def translate_exception(exception, message:, sql:, binds:)
+          if exception.is_a?(::Mysql2::Error::TimeoutError) && !exception.error_number
+            ActiveRecord::AdapterTimeout.new(message, sql: sql, binds: binds, connection_pool: @pool)
+          elsif exception.is_a?(::Mysql2::Error::ConnectionError)
+            if exception.message.match?(/MySQL client is not connected/i)
+              ActiveRecord::ConnectionNotEstablished.new(exception, connection_pool: @pool)
+            else
+              ActiveRecord::ConnectionFailed.new(message, sql: sql, binds: binds, connection_pool: @pool)
+            end
+          else
+            super
+          end
+        end
+
+        def default_prepared_statements
+          false
+        end
+
+        ActiveRecord::Type.register(:immutable_string, adapter: :mysql2) do |_, **args|
+          Type::ImmutableString.new(true: "1", false: "0", **args)
+        end
+
+        ActiveRecord::Type.register(:string, adapter: :mysql2) do |_, **args|
+          Type::String.new(true: "1", false: "0", **args)
+        end
+
+        ActiveRecord::Type.register(:unsigned_integer, Type::UnsignedInteger, adapter: :mysql2)
+    end
+
+    ActiveSupport.run_load_hooks(:active_record_mysql2adapter, Mysql2Adapter)
+  end
+end

data/lib/active_record/connection_adapters/pool_config.rb

@@ -0,0 +1,83 @@
+# frozen_string_literal: true
+
+module ActiveRecord
+  module ConnectionAdapters
+    class PoolConfig # :nodoc:
+      include MonitorMixin
+
+      attr_reader :db_config, :role, :shard
+      attr_writer :schema_reflection, :server_version
+      attr_accessor :connection_class
+
+      def schema_reflection
+        @schema_reflection ||= SchemaReflection.new(db_config.lazy_schema_cache_path)
+      end
+
+      INSTANCES = ObjectSpace::WeakMap.new
+      private_constant :INSTANCES
+
+      class << self
+        def discard_pools!
+          INSTANCES.each_key(&:discard_pool!)
+        end
+
+        def disconnect_all!
+          INSTANCES.each_key { |c| c.disconnect!(automatic_reconnect: true) }
+        end
+      end
+
+      def initialize(connection_class, db_config, role, shard)
+        super()
+        @server_version = nil
+        @connection_class = connection_class
+        @db_config = db_config
+        @role = role
+        @shard = shard
+        @pool = nil
+        INSTANCES[self] = self
+      end
+
+      def server_version(connection)
+        @server_version || synchronize { @server_version ||= connection.get_database_version }
+      end
+
+      def connection_name
+        if connection_class.primary_class?
+          "ActiveRecord::Base"
+        else
+          connection_class.name
+        end
+      end
+
+      def disconnect!(automatic_reconnect: false)
+        return unless @pool
+
+        synchronize do
+          return unless @pool
+
+          @pool.automatic_reconnect = automatic_reconnect
+          @pool.disconnect!
+        end
+
+        nil
+      end
+
+      def pool
+        @pool || synchronize { @pool ||= ConnectionAdapters::ConnectionPool.new(self) }
+      end
+
+      def discard_pool!
+        return unless @pool
+
+        synchronize do
+          return unless @pool
+
+          @pool.discard!
+          @pool = nil
+        end
+      end
+    end
+  end
+end
+
+ActiveSupport::ForkTracker.after_fork { ActiveRecord::ConnectionAdapters::PoolConfig.discard_pools! }

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

@@ -0,0 +1,82 @@
+# frozen_string_literal: true
+
+module ActiveRecord
+  module ConnectionAdapters
+    module PostgreSQL
+      class Column < ConnectionAdapters::Column # :nodoc:
+        delegate :oid, :fmod, to: :sql_type_metadata
+
+        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.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
+    PostgreSQLColumn = PostgreSQL::Column # :nodoc:
+  end
+end

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

@@ -0,0 +1,231 @@
+# frozen_string_literal: true
+
+module ActiveRecord
+  module ConnectionAdapters
+    module PostgreSQL
+      module DatabaseStatements
+        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:
+          result = internal_execute(sql, name)
+          result.map_types!(@type_map_for_results).values
+        end
+
+        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(...) # :nodoc:
+          super
+        ensure
+          @notice_receiver_sql_warnings = []
+        end
+
+        def exec_insert(sql, name = nil, binds = [], pk = nil, sequence_name = nil, returning: nil) # :nodoc:
+          if use_insert_returning? || pk == false
+            super
+          else
+            result = internal_exec_query(sql, name, binds)
+            unless sequence_name
+              table_ref = extract_table_ref_from_insert_sql(sql)
+              if table_ref
+                pk = primary_key(table_ref) if pk.nil?
+                pk = suppress_composite_primary_key(pk)
+                sequence_name = default_sequence_name(table_ref, pk)
+              end
+              return result unless sequence_name
+            end
+            last_insert_id_result(sequence_name)
+          end
+        end
+
+        # Begins a transaction.
+        def begin_db_transaction # :nodoc:
+          internal_execute("BEGIN", "TRANSACTION", allow_retry: true, materialize_transactions: false)
+        end
+
+        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 # :nodoc:
+          internal_execute("COMMIT", "TRANSACTION", allow_retry: false, materialize_transactions: true)
+        end
+
+        # Aborts a transaction.
+        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
+          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 perform_query(raw_connection, sql, binds, type_casted_binds, prepare:, notification_payload:, batch: false)
+            update_typemap_for_default_timezone
+            result = if prepare
+              begin
+                stmt_key = prepare_statement(sql, binds, raw_connection)
+                notification_payload[:statement_name] = stmt_key
+                raw_connection.exec_prepared(stmt_key, type_casted_binds)
+              rescue PG::FeatureNotSupported => error
+                if is_cached_plan_failure?(error)
+                  # Nothing we can do if we are in a transaction because all commands
+                  # will raise InFailedSQLTransaction
+                  if in_transaction?
+                    raise PreparedStatementCacheExpired.new(error.message, connection_pool: @pool)
+                  else
+                    @lock.synchronize do
+                      # outside of transactions we can simply flush this query and retry
+                      @statements.delete sql_key(sql)
+                    end
+                    retry
+                  end
+                end
+
+                raise
+              end
+            elsif binds.nil? || binds.empty?
+              raw_connection.async_exec(sql)
+            else
+              raw_connection.exec_params(sql, type_casted_binds)
+            end
+
+            verified!
+            handle_warnings(result)
+            notification_payload[:row_count] = result.count
+            result
+          end
+
+          def cast_result(result)
+            if result.fields.empty?
+              result.clear
+              return ActiveRecord::Result.empty
+            end
+
+            types = {}
+            fields = result.fields
+            fields.each_with_index do |fname, i|
+              ftype = result.ftype i
+              fmod  = result.fmod i
+              types[fname] = types[i] = get_oid_type(ftype, fmod, fname)
+            end
+            ar_result = ActiveRecord::Result.new(fields, result.values, types.freeze)
+            result.clear
+            ar_result
+          end
+
+          def affected_rows(result)
+            affected_rows = result.cmd_tuples
+            result.clear
+            affected_rows
+          end
+
+          def execute_batch(statements, name = nil, **kwargs)
+            raw_execute(combine_multi_statements(statements), name, batch: true, **kwargs)
+          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)
+            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
+end

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

@@ -0,0 +1,44 @@
+# frozen_string_literal: true
+
+module ActiveRecord
+  module ConnectionAdapters
+    module PostgreSQL
+      class ExplainPrettyPrinter # :nodoc:
+        # Pretty prints the result of an EXPLAIN in a way that resembles the output of the
+        # PostgreSQL shell:
+        #
+        #                                     QUERY PLAN
+        #   ------------------------------------------------------------------------------
+        #    Nested Loop Left Join  (cost=0.00..37.24 rows=8 width=0)
+        #      Join Filter: (posts.user_id = users.id)
+        #      ->  Index Scan using users_pkey on users  (cost=0.00..8.27 rows=1 width=4)
+        #            Index Cond: (id = 1)
+        #      ->  Seq Scan on posts  (cost=0.00..28.88 rows=8 width=4)
+        #            Filter: (posts.user_id = 1)
+        #   (6 rows)
+        #
+        def pp(result)
+          header = result.columns.first
+          lines  = result.rows.map(&:first)
+
+          # We add 2 because there's one char of padding at both sides, note
+          # the extra hyphens in the example above.
+          width = [header, *lines].map(&:length).max + 2
+
+          pp = []
+
+          pp << header.center(width).rstrip
+          pp << "-" * width
+
+          pp += lines.map { |line| " #{line}" }
+
+          nrows = result.rows.length
+          rows_label = nrows == 1 ? "row" : "rows"
+          pp << "(#{nrows} #{rows_label})"
+
+          pp.join("\n") + "\n"
+        end
+      end
+    end
+  end
+end