Feature enrichment for XAResourceWrapper

agroal-narayana/src/main/java/io/agroal/narayana/FirstResourceBaseXAResource.java

@@ -0,0 +1,16 @@
+// Copyright (C) 2020 Red Hat, Inc. and individual contributors as indicated by the @author tags.
+// You may not use this file except in compliance with the Apache License, Version 2.0.
+
+package io.agroal.narayana;
+
+import io.agroal.api.transaction.TransactionAware;
+import org.jboss.tm.FirstResource;
+
+import javax.transaction.xa.XAResource;
+
+public class FirstResourceBaseXAResource extends BaseXAResource implements FirstResource {
+
+    public FirstResourceBaseXAResource(TransactionAware transactionAware, XAResource xaResource, String jndiName) {
+        super( transactionAware, xaResource, jndiName );
+    }
+}

agroal-narayana/src/main/java/io/agroal/narayana/LocalXAResource.java

@@ -4,6 +4,7 @@
 package io.agroal.narayana;
 
 import io.agroal.api.transaction.TransactionAware;
+import org.jboss.tm.LastResource;
 import org.jboss.tm.XAResourceWrapper;
 
 import javax.transaction.xa.XAException;
@@ -19,11 +20,11 @@
  * @author <a href="lbarreiro@redhat.com">Luis Barreiro</a>
  * @author <a href="jesper.pedersen@redhat.com">Jesper Pedersen</a>
  */
-public class LocalXAResource implements XAResourceWrapper {
+public class LocalXAResource implements XAResourceWrapper, LastResource {
 
     private static final String PRODUCT_NAME = LocalXAResource.class.getPackage().getImplementationTitle();
     private static final String PRODUCT_VERSION = LocalXAResource.class.getPackage().getImplementationVersion();
-    
+
     private final TransactionAware transactionAware;
     private final String jndiName;
     private Xid currentXid;

agroal-narayana/src/main/java/io/agroal/narayana/NarayanaTransactionIntegration.java

@@ -42,6 +42,8 @@ public class NarayanaTransactionIntegration implements TransactionIntegration {
 
     private final boolean connectable;
 
+    private final boolean firstResource;
+
     private final XAResourceRecoveryRegistry recoveryRegistry;
 
     // In order to construct a UID that is globally unique, simply pair a UID with an InetAddress.
@@ -60,10 +62,22 @@ public NarayanaTransactionIntegration(TransactionManager transactionManager, Tra
     }
 
     public NarayanaTransactionIntegration(TransactionManager transactionManager, TransactionSynchronizationRegistry transactionSynchronizationRegistry, String jndiName, boolean connectable, XAResourceRecoveryRegistry recoveryRegistry) {
+        this( transactionManager, transactionSynchronizationRegistry, jndiName, connectable, false, recoveryRegistry );
+    }
+
+    public NarayanaTransactionIntegration(TransactionManager transactionManager, TransactionSynchronizationRegistry transactionSynchronizationRegistry, String jndiName, boolean connectable, boolean firstResource) {
+        this( transactionManager, transactionSynchronizationRegistry, jndiName, connectable, firstResource, null );
+    }
+
+    public NarayanaTransactionIntegration(TransactionManager transactionManager, TransactionSynchronizationRegistry transactionSynchronizationRegistry, String jndiName, boolean connectable, boolean firstResource, XAResourceRecoveryRegistry recoveryRegistry) {
+        if (connectable && firstResource) {
+            throw new IllegalArgumentException("Setting connectable and firstResource to true is disallowed at the same time.");
+        }
         this.transactionManager = transactionManager;
         this.transactionSynchronizationRegistry = transactionSynchronizationRegistry;
         this.jndiName = jndiName;
         this.connectable = connectable;
+        this.firstResource = firstResource;
         this.recoveryRegistry = recoveryRegistry;
     }
 
@@ -82,6 +96,20 @@ public enum TransactionPhase {
         TRANSACTION_NONE, TRANSACTION_ACTIVE, TRANSACTION_COMPLETING, TRANSACTION_DONE
     }
 
+    private XAResource createXaResource(TransactionAware transactionAware, XAResource xaResource) {
+        if ( xaResource != null ) {
+            if ( firstResource ) {
+                return new FirstResourceBaseXAResource( transactionAware, xaResource, jndiName );
+            } else {
+                return new BaseXAResource( transactionAware, xaResource, jndiName );
+            }
+        } else if ( connectable ) {
+            return new ConnectableLocalXAResource( transactionAware, jndiName );
+        } else {
+            return new LocalXAResource( transactionAware, jndiName );
+        }
+    }
+
     @Override
     public void associate(TransactionAware transactionAware, XAResource xaResource) throws SQLException {
         try {
@@ -90,16 +118,7 @@ public void associate(TransactionAware transactionAware, XAResource xaResource)
                 if ( transactionSynchronizationRegistry.getResource( key ) == null ) {
                     transactionSynchronizationRegistry.registerInterposedSynchronization( new InterposedSynchronization( transactionAware ) );
                     transactionSynchronizationRegistry.putResource( key, transactionAware );
-
-                    XAResource xaResourceToEnlist;
-                    if ( xaResource != null ) {
-                        xaResourceToEnlist = new BaseXAResource( transactionAware, xaResource, jndiName );
-                    } else if ( connectable ) {
-                        xaResourceToEnlist = new ConnectableLocalXAResource( transactionAware, jndiName );
-                    } else {
-                        xaResourceToEnlist = new LocalXAResource( transactionAware, jndiName );
-                    }
-                    transactionManager.getTransaction().enlistResource( xaResourceToEnlist );
+                    transactionManager.getTransaction().enlistResource( createXaResource( transactionAware, xaResource ) );
                 } else {
                     transactionAware.transactionStart();
                 }

agroal-spring-boot-starter/src/main/java/io/agroal/springframework/boot/AgroalDataSourceConfiguration.java

@@ -40,7 +40,10 @@ public class AgroalDataSourceConfiguration {
 
     @Bean
     @ConfigurationProperties( prefix = "spring.datasource.agroal" )
-    public AgroalDataSource dataSource(DataSourceProperties properties, @Value("${spring.datasource.agroal.connectable:false}") boolean connectable) {
+    public AgroalDataSource dataSource(
+            DataSourceProperties properties,
+            @Value("${spring.datasource.agroal.connectable:false}") boolean connectable,
+            @Value("${spring.datasource.agroal.firstResource:false}") boolean firstResource) {
         AgroalDataSource dataSource = properties.initializeDataSourceBuilder().type( AgroalDataSource.class ).build();
         if ( !StringUtils.hasLength( properties.getDriverClassName() ) ) {
             DatabaseDriver driver = DatabaseDriver.fromJdbcUrl( properties.determineUrl() );
@@ -58,6 +61,7 @@ public AgroalDataSource dataSource(DataSourceProperties properties, @Value("${sp
                     jtaPlatform.getTransactionSynchronizationRegistry(),
                     name,
                     connectable,
+                    firstResource,
                     recoveryRegistry );
             dataSource.setJtaTransactionIntegration( transactionIntegration );
         }

agroal-test/src/test/java/io/agroal/test/narayana/EnlistmentTests.java

@@ -8,10 +8,10 @@
 import io.agroal.api.configuration.supplier.AgroalDataSourceConfigurationSupplier;
 import io.agroal.narayana.BaseXAResource;
 import io.agroal.narayana.ConnectableLocalXAResource;
+import io.agroal.narayana.FirstResourceBaseXAResource;
 import io.agroal.narayana.LocalXAResource;
 import io.agroal.narayana.NarayanaTransactionIntegration;
 import io.agroal.test.MockXADataSource;
-import jakarta.transaction.Synchronization;
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.DisplayName;
@@ -23,6 +23,7 @@
 import jakarta.transaction.InvalidTransactionException;
 import jakarta.transaction.NotSupportedException;
 import jakarta.transaction.RollbackException;
+import jakarta.transaction.Synchronization;
 import jakarta.transaction.SystemException;
 import jakarta.transaction.Transaction;
 import jakarta.transaction.TransactionManager;
@@ -232,29 +233,40 @@ void enlistedResourceTypeTest() throws SQLException, SystemException {
         AgroalDataSourceConfiguration regularConfiguration = new AgroalDataSourceConfigurationSupplier()
                 .connectionPoolConfiguration( cp -> cp
                         .maxSize( 1 )
-                        .transactionIntegration( new NarayanaTransactionIntegration( txManager, txSyncRegistry, null, false ) )
+                        .transactionIntegration( new NarayanaTransactionIntegration( txManager, txSyncRegistry, null, false, false ) )
                 ).get();
 
         AgroalDataSourceConfiguration connectableConfiguration = new AgroalDataSourceConfigurationSupplier()
                 .connectionPoolConfiguration( cp -> cp
                         .maxSize( 1 )
-                        .transactionIntegration( new NarayanaTransactionIntegration( txManager, txSyncRegistry, null, true ) )
+                        .transactionIntegration( new NarayanaTransactionIntegration( txManager, txSyncRegistry, null, true, false ) )
                 ).get();
 
         AgroalDataSourceConfiguration xaConfiguration = new AgroalDataSourceConfigurationSupplier()
                 .connectionPoolConfiguration( cp -> cp
                         .maxSize( 1 )
-                        .transactionIntegration( new NarayanaTransactionIntegration( txManager, txSyncRegistry, null, true ) )
+                        .transactionIntegration( new NarayanaTransactionIntegration( txManager, txSyncRegistry, null, false, false ) )
                         .connectionFactoryConfiguration( cf -> cf
                                 .connectionProviderClass( MockXADataSource.Empty.class )
                         )
                 ).get();
 
+        AgroalDataSourceConfiguration firstXaConfiguration = new AgroalDataSourceConfigurationSupplier()
+                .connectionPoolConfiguration( cp -> cp
+                        .maxSize( 1 )
+                        .transactionIntegration( new NarayanaTransactionIntegration( txManager, txSyncRegistry, null, false, true ) )
+                        .connectionFactoryConfiguration( cf -> cf
+                                .connectionProviderClass( MockXADataSource.Empty.class )
+                        )
+                ).get();
+
         verifyEnlistedResourceType( regularConfiguration, txManager, "regular", LocalXAResource.class );
 
         verifyEnlistedResourceType( connectableConfiguration, txManager, "connectable", ConnectableLocalXAResource.class );
 
         verifyEnlistedResourceType( xaConfiguration, txManager, "xa", BaseXAResource.class );
+
+        verifyEnlistedResourceType( firstXaConfiguration, txManager, "firstResource xa", FirstResourceBaseXAResource.class );
     }
 
     private static void verifyEnlistedResourceType(AgroalDataSourceConfiguration configuration, TransactionManager txManager, String type, Class<?> resourceClass) throws SQLException, SystemException {