Refactor PostgresqlAdapter#raw_execute to be less unique

The postgres adapter used to be more complex than the others to
be able to pass the prepared statement key to the `log` method.

If we instead add it to the payload later, we can simplify the method
further and it opens the door to refactor `log` and `with_raw_connection`
out of `raw_execute`.

Author: byroot

activerecord/lib/active_record/connection_adapters/abstract_adapter.rb

@@ -1118,14 +1118,13 @@ def translate_exception_class(native_error, sql, binds)
           active_record_error
         end
 
-        def log(sql, name = "SQL", binds = [], type_casted_binds = [], statement_name = nil, async: false, &block) # :doc:
+        def log(sql, name = "SQL", binds = [], type_casted_binds = [], async: false, &block) # :doc:
           @instrumenter.instrument(
             "sql.active_record",
             sql:               sql,
             name:              name,
             binds:             binds,
             type_casted_binds: type_casted_binds,
-            statement_name:    statement_name,
             async:             async,
             connection:        self,
             transaction:       current_transaction.user_transaction.presence,

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

@@ -137,49 +137,40 @@ def raw_execute(sql, name, binds = nil, prepare: false, async: false, allow_retr
 
             type_casted_binds = type_casted_binds(binds)
 
-            begin
+            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|
-                stmt_key = prepare_statement(sql, binds, conn) if prepare
-
-                log(sql, name, binds, type_casted_binds, stmt_key, async: async) do |notification_payload|
+                result = if prepare
                   begin
-                    result = if prepare
-                      conn.exec_prepared(stmt_key, type_casted_binds)
-                    elsif without_prepared_statement?(binds)
-                      conn.async_exec(sql)
-                    else
-                      conn.exec_params(sql, type_casted_binds)
+                    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
                     end
-                  rescue => original_exception
-                    # Contrary to all other adapters we have to enter `with_raw_connection` before `log`
-                    # so that we can prepare the statement and pass the key to `log`.
-                    # So we need to translate exceptions ourselves.
-                    raise translate_exception_class(original_exception, sql, binds)
-                  end
 
-                  verified!
-                  handle_warnings(result)
-                  notification_payload[:row_count] = result.count
-                  result
-                end
-              end
-            rescue ActiveRecord::StatementInvalid => error
-              if prepare
-                raise unless 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 ActiveRecord::PreparedStatementCacheExpired.new(error.cause.message, connection_pool: @pool)
-                else
-                  @lock.synchronize do
-                    # outside of transactions we can simply flush this query and retry
-                    @statements.delete sql_key(sql)
+                    raise
                   end
-                  retry
+                elsif without_prepared_statement?(binds)
+                  conn.async_exec(sql)
+                else
+                  conn.exec_params(sql, type_casted_binds)
                 end
-              else
-                raise
+
+                verified!
+                handle_warnings(result)
+                notification_payload[:row_count] = result.count
+                result
               end
             end
           end

activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb

@@ -872,8 +872,7 @@ def load_types_queries(initializer, oids)
         #
         # Check here for more details:
         # https://git.postgresql.org/gitweb/?p=postgresql.git;a=blob;f=src/backend/utils/cache/plancache.c#l573
-        def is_cached_plan_failure?(e)
-          pgerror = e.cause
+        def is_cached_plan_failure?(pgerror)
           pgerror.result.result_error_field(PG::PG_DIAG_SQLSTATE) == FEATURE_NOT_SUPPORTED &&
             pgerror.result.result_error_field(PG::PG_DIAG_SOURCE_FUNCTION) == "RevalidateCachedQuery"
         rescue

activerecord/test/cases/adapters/postgresql/connection_test.rb

@@ -130,8 +130,11 @@ def test_schema_names_logs_name
       def test_statement_key_is_logged
         bind = Relation::QueryAttribute.new(nil, 1, Type::Value.new)
         @connection.exec_query("SELECT $1::integer", "SQL", [bind], prepare: true)
-        name = @subscriber.payloads.last[:statement_name]
-        assert name
+
+        payload = @subscriber.payloads.find { |p| p[:sql] == "SELECT $1::integer" }
+        name = payload[:statement_name]
+        assert_not_nil name
+
         res = @connection.exec_query("EXPLAIN (FORMAT JSON) EXECUTE #{name}(1)")
         plan = res.column_types["QUERY PLAN"].deserialize res.rows.first.first
         assert_operator plan.length, :>, 0

activerecord/test/cases/adapters/trilogy/trilogy_adapter_test.rb

@@ -222,8 +222,6 @@ class TrilogyAdapterTest < ActiveRecord::TrilogyTestCase
         assert_includes payload, :type_casted_binds
         assert_equal [], payload[:type_casted_binds]
 
-        # :stament_name is always nil and never set 🤷‍♂️
-        assert_includes payload, :statement_name
         assert_nil payload[:statement_name]
 
         assert_not_includes payload, :cached