Refactor Active Record connection adapters further

A `raw_execute` implementation is now provided, instead adapters
have to implement `perform_query`.

It's a much simpler method that no longer need to concern itself
with Active Support notifications nor calling `with_raw_connection`.

Author: byroot

activerecord/lib/active_record/connection_adapters/abstract/database_statements.rb

@@ -545,6 +545,28 @@ def internal_exec_query(...) # :nodoc:
       private
         # Lowest level way to execute a query. Doesn't check for illegal writes, doesn't annotate queries, yields a native result object.
         def raw_execute(sql, name = nil, binds = [], prepare: false, async: false, allow_retry: false, materialize_transactions: true)
+          type_casted_binds = type_casted_binds(binds)
+          notification_payload = {
+            sql: sql,
+            name: name,
+            binds: binds,
+            type_casted_binds: type_casted_binds,
+            async: async,
+            connection: self,
+            transaction: current_transaction.user_transaction.presence,
+            statement_name: nil,
+            row_count: 0,
+          }
+          @instrumenter.instrument("sql.active_record", notification_payload) do
+            with_raw_connection(allow_retry: allow_retry, materialize_transactions: materialize_transactions) do |conn|
+              perform_query(conn, sql, binds, type_casted_binds, prepare: prepare, notification_payload: notification_payload)
+            end
+          rescue ActiveRecord::StatementInvalid => ex
+            raise ex.set_query(sql, binds)
+          end
+        end
+
+        def perform_query(raw_connection, sql, binds, type_casted_binds, prepare:, notification_payload:)
           raise NotImplementedError
         end
 

activerecord/lib/active_record/connection_adapters/abstract_mysql_adapter.rb

@@ -773,6 +773,7 @@ def warning_ignored?(warning)
         # Make sure we carry over any changes to ActiveRecord.default_timezone that have been
         # made since we established the connection
         def sync_timezone_changes(raw_connection)
+          raise NotImplementedError
         end
 
         # See https://dev.mysql.com/doc/mysql-errors/en/server-error-reference.html

activerecord/lib/active_record/connection_adapters/mysql2/database_statements.rb

@@ -58,37 +58,34 @@ def with_multi_statements
             end
           end
 
-          def raw_execute(sql, name, binds = nil, prepare: false, 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|
-                sync_timezone_changes(conn)
-
-                result = if prepare
-                  stmt = @statements[sql] ||= conn.prepare(sql)
-
-                  begin
-                    ActiveSupport::Dependencies.interlock.permit_concurrent_loads do
-                      stmt.execute(*type_casted_binds)
-                    end
-                  rescue ::Mysql2::Error
-                    @statements.delete(sql)
-                    stmt.close
-                    raise
-                  end
-                  verified!
-                else
-                  conn.query(sql)
-                end
+          def perform_query(raw_connection, sql, binds, type_casted_binds, prepare:, notification_payload:)
+            sync_timezone_changes(raw_connection)
 
-                @affected_rows_before_warnings = conn.affected_rows
-                conn.abandon_results!
+            result = if prepare
+              stmt = @statements[sql] ||= raw_connection.prepare(sql)
 
-                verified!
-                handle_warnings(sql)
-                notification_payload[:row_count] = result&.size || 0
-                result
+              begin
+                ActiveSupport::Dependencies.interlock.permit_concurrent_loads do
+                  stmt.execute(*type_casted_binds)
+                end
+              rescue ::Mysql2::Error
+                @statements.delete(sql)
+                stmt.close
+                raise
               end
+              verified!
+            else
+              raw_connection.query(sql)
             end
+
+            notification_payload[:row_count] = result&.size || 0
+
+            @affected_rows_before_warnings = raw_connection.affected_rows
+            raw_connection.abandon_results!
+
+            verified!
+            handle_warnings(sql)
+            result
           end
 
           def cast_result(result)

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

@@ -132,47 +132,40 @@ def cancel_any_running_query
           rescue PG::Error
           end
 
-          def raw_execute(sql, name, binds = nil, prepare: false, async: false, allow_retry: false, materialize_transactions: true)
+          def perform_query(raw_connection, sql, binds, type_casted_binds, prepare:, notification_payload:)
             update_typemap_for_default_timezone
-
-            type_casted_binds = type_casted_binds(binds)
-
-            log(sql, name, binds, type_casted_binds, async: async) do |notification_payload|
-              with_raw_connection(allow_retry: allow_retry, materialize_transactions: materialize_transactions) do |conn|
-                result = if prepare
-                  begin
-                    stmt_key = prepare_statement(sql, binds, conn)
-                    notification_payload[:statement_name] = stmt_key
-                    conn.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
+            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
-
-                    raise
+                    retry
                   end
-                elsif without_prepared_statement?(binds)
-                  conn.async_exec(sql)
-                else
-                  conn.exec_params(sql, type_casted_binds)
                 end
 
-                verified!
-                handle_warnings(result)
-                notification_payload[:row_count] = result.count
-                result
+                raise
               end
+            elsif without_prepared_statement?(binds)
+              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)

activerecord/lib/active_record/connection_adapters/sqlite3/database_statements.rb

@@ -86,46 +86,40 @@ def internal_begin_transaction(mode, isolation)
             end
           end
 
-          def raw_execute(sql, name = nil, binds = [], prepare: false, async: false, allow_retry: false, materialize_transactions: true) # :nodoc:
-            type_casted_binds = type_casted_binds(binds)
-
-            result = log(sql, name, binds, type_casted_binds, async: async) do |notification_payload|
-              with_raw_connection(materialize_transactions: materialize_transactions) do |conn|
-                if prepare
-                  stmt = @statements[sql] ||= conn.prepare(sql)
-                  stmt.reset!
+          def perform_query(raw_connection, sql, binds, type_casted_binds, prepare:, notification_payload:)
+            if prepare
+              stmt = @statements[sql] ||= raw_connection.prepare(sql)
+              stmt.reset!
+              stmt.bind_params(type_casted_binds)
+
+              result = if stmt.column_count.zero? # No return
+                stmt.step
+                ActiveRecord::Result.empty
+              else
+                ActiveRecord::Result.new(stmt.columns, stmt.to_a)
+              end
+            else
+              # Don't cache statements if they are not prepared.
+              stmt = raw_connection.prepare(sql)
+              begin
+                unless without_prepared_statement?(binds)
                   stmt.bind_params(type_casted_binds)
-
-                  result = if stmt.column_count.zero? # No return
-                    stmt.step
-                    ActiveRecord::Result.empty
-                  else
-                    ActiveRecord::Result.new(stmt.columns, stmt.to_a)
-                  end
+                end
+                result = if stmt.column_count.zero? # No return
+                  stmt.step
+                  ActiveRecord::Result.empty
                 else
-                  # Don't cache statements if they are not prepared.
-                  stmt = conn.prepare(sql)
-                  begin
-                    unless without_prepared_statement?(binds)
-                      stmt.bind_params(type_casted_binds)
-                    end
-                    result = if stmt.column_count.zero? # No return
-                      stmt.step
-                      ActiveRecord::Result.empty
-                    else
-                      ActiveRecord::Result.new(stmt.columns, stmt.to_a)
-                    end
-                  ensure
-                    stmt.close
-                  end
+                  ActiveRecord::Result.new(stmt.columns, stmt.to_a)
                 end
-                @last_affected_rows = @raw_connection.changes
-                verified!
-
-                notification_payload[:row_count] = result.length
-                result
+              ensure
+                stmt.close
               end
             end
+            @last_affected_rows = raw_connection.changes
+            verified!
+
+            notification_payload[:row_count] = result.length
+            result
           end
 
           def cast_result(result)

activerecord/lib/active_record/connection_adapters/trilogy/database_statements.rb

@@ -10,20 +10,16 @@ def exec_insert(sql, name, binds, pk = nil, sequence_name = nil, returning: nil)
         end
 
         private
-          def raw_execute(sql, name, binds = nil, prepare: false, 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|
-                sync_timezone_changes(conn)
-                result = conn.query(sql)
-                while conn.more_results_exist?
-                  conn.next_result
-                end
-                verified!
-                handle_warnings(sql)
-                notification_payload[:row_count] = result.count
-                result
-              end
+          def perform_query(raw_connection, sql, binds, type_casted_binds, prepare:, notification_payload:)
+            sync_timezone_changes(raw_connection)
+            result = raw_connection.query(sql)
+            while raw_connection.more_results_exist?
+              raw_connection.next_result
             end
+            verified!
+            handle_warnings(sql)
+            notification_payload[:row_count] = result.count
+            result
           end
 
           def cast_result(result)