Skip to content

Commit

Permalink
[WFLY-18016] Refactoring part of AGROAL subsystem - eliminating usage…
Browse files Browse the repository at this point in the history
… of deprecated CapabilityServiceBuilder.addCapabilityRequirement() methods
  • Loading branch information
ropalka committed Jun 9, 2023
1 parent 61baae8 commit 2feaab5
Show file tree
Hide file tree
Showing 4 changed files with 64 additions and 54 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,7 @@
import java.sql.Connection;
import java.sql.SQLException;
import java.time.Duration;
import java.util.function.Supplier;

import static io.agroal.api.configuration.AgroalConnectionPoolConfiguration.ConnectionValidator.defaultValidator;
import static java.time.Duration.ofMillis;
Expand Down Expand Up @@ -118,16 +119,21 @@ protected static AgroalConnectionFactoryConfigurationSupplier connectionFactoryC
return configuration;
}

protected static void setupElytronSecurity(OperationContext context, ModelNode model, DataSourceService dataSourceService, ServiceBuilder<?> serviceBuilder) throws OperationFailedException {
protected static Supplier<AuthenticationContext> setupAuthenticationContext(OperationContext context, ModelNode model, ServiceBuilder<?> serviceBuilder) throws OperationFailedException {
if (AbstractDataSourceDefinition.AUTHENTICATION_CONTEXT.resolveModelAttribute(context, model).isDefined()) {
String authenticationContextName = AbstractDataSourceDefinition.AUTHENTICATION_CONTEXT.resolveModelAttribute(context, model).asString();
ServiceName authenticationContextCapability = context.getCapabilityServiceName(AbstractDataSourceDefinition.AUTHENTICATION_CONTEXT_CAPABILITY, authenticationContextName, AuthenticationContext.class);
serviceBuilder.addDependency(authenticationContextCapability, AuthenticationContext.class, dataSourceService.getAuthenticationContextInjector());
return serviceBuilder.requires(authenticationContextCapability);
}
return null;
}

protected static Supplier<ExceptionSupplier<CredentialSource, Exception>> setupCredentialReference(OperationContext context, ModelNode model, ServiceBuilder<?> serviceBuilder) throws OperationFailedException {
if (AbstractDataSourceDefinition.CREDENTIAL_REFERENCE.resolveModelAttribute(context, model).isDefined()) {
ExceptionSupplier<CredentialSource, Exception> credentialSourceSupplier = CredentialReference.getCredentialSourceSupplier(context, AbstractDataSourceDefinition.CREDENTIAL_REFERENCE, model, serviceBuilder);
dataSourceService.getCredentialSourceSupplierInjector().inject(credentialSourceSupplier);
return () -> credentialSourceSupplier;
}
return null;
}

protected static AgroalConnectionPoolConfigurationSupplier connectionPoolConfiguration(OperationContext context, ModelNode model) throws OperationFailedException {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -25,11 +25,14 @@
import static org.jboss.as.controller.security.CredentialReference.rollbackCredentialStoreUpdate;
import static org.wildfly.extension.datasources.agroal.AbstractDataSourceDefinition.CREDENTIAL_REFERENCE;

import jakarta.transaction.TransactionSynchronizationRegistry;
import java.util.function.Consumer;
import java.util.function.Supplier;

import io.agroal.api.AgroalDataSource;
import io.agroal.api.configuration.supplier.AgroalConnectionFactoryConfigurationSupplier;
import io.agroal.api.configuration.supplier.AgroalConnectionPoolConfigurationSupplier;
import io.agroal.api.configuration.supplier.AgroalDataSourceConfigurationSupplier;
import jakarta.transaction.TransactionSynchronizationRegistry;
import org.jboss.as.controller.AbstractAddStepHandler;
import org.jboss.as.controller.AbstractRemoveStepHandler;
import org.jboss.as.controller.CapabilityServiceBuilder;
Expand All @@ -39,6 +42,9 @@
import org.jboss.as.controller.registry.Resource;
import org.jboss.dmr.ModelNode;
import org.jboss.msc.service.ServiceName;
import org.wildfly.common.function.ExceptionSupplier;
import org.wildfly.security.auth.client.AuthenticationContext;
import org.wildfly.security.credential.source.CredentialSource;

/**
* Operations for adding and removing a datasource resource to the model
Expand Down Expand Up @@ -89,18 +95,15 @@ protected void performRuntime(OperationContext context, ModelNode operation, Mod
boolean jta = DataSourceDefinition.JTA_ATTRIBUTE.resolveModelAttribute(context, model).asBoolean();
boolean connectable = DataSourceDefinition.CONNECTABLE_ATTRIBUTE.resolveModelAttribute(context, model).asBoolean();
String driverName = AbstractDataSourceDefinition.DRIVER_ATTRIBUTE.resolveModelAttribute(context, factoryModel).asString();

DataSourceService dataSourceService = new DataSourceService(datasourceName, jndiName, jta, connectable, false, dataSourceConfiguration);

CapabilityServiceBuilder serviceBuilder = context.getCapabilityServiceTarget().addCapability(AbstractDataSourceDefinition.DATA_SOURCE_CAPABILITY.fromBaseCapability(datasourceName))
.setInstance(dataSourceService)
.addCapabilityRequirement(DriverDefinition.AGROAL_DRIVER_CAPABILITY.getDynamicName(driverName), Class.class, dataSourceService.getDriverInjector());
if (jta) {
// TODO add a Stage.MODEL requirement
serviceBuilder.addCapabilityRequirement("org.wildfly.transactions.transaction-synchronization-registry", TransactionSynchronizationRegistry.class, dataSourceService.getTransactionSynchronizationRegistryInjector());
}
AbstractDataSourceOperations.setupElytronSecurity(context, factoryModel, dataSourceService, serviceBuilder);

final CapabilityServiceBuilder serviceBuilder = context.getCapabilityServiceTarget().addCapability(AbstractDataSourceDefinition.DATA_SOURCE_CAPABILITY.fromBaseCapability(datasourceName));
final Consumer<AgroalDataSource> consumer = serviceBuilder.provides(AbstractDataSourceDefinition.DATA_SOURCE_CAPABILITY.fromBaseCapability(datasourceName));
final Supplier<Class> driverSupplier = serviceBuilder.requiresCapability(DriverDefinition.AGROAL_DRIVER_CAPABILITY.getDynamicName(driverName), Class.class);
final Supplier<AuthenticationContext> authenticationContextSupplier = AbstractDataSourceOperations.setupAuthenticationContext(context, factoryModel, serviceBuilder);
final Supplier<ExceptionSupplier<CredentialSource, Exception>> credentialSourceSupplier = AbstractDataSourceOperations.setupCredentialReference(context, factoryModel, serviceBuilder);
// TODO add a Stage.MODEL requirement
final Supplier<TransactionSynchronizationRegistry> txnRegistrySupplier = jta ? serviceBuilder.requiresCapability("org.wildfly.transactions.transaction-synchronization-registry", TransactionSynchronizationRegistry.class) : null;
DataSourceService dataSourceService = new DataSourceService(consumer, driverSupplier, authenticationContextSupplier, credentialSourceSupplier, txnRegistrySupplier, datasourceName, jndiName, jta, connectable, false, dataSourceConfiguration);
serviceBuilder.setInstance(dataSourceService);
serviceBuilder.install();
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,6 @@
import org.jboss.msc.service.StartContext;
import org.jboss.msc.service.StartException;
import org.jboss.msc.service.StopContext;
import org.jboss.msc.value.InjectedValue;
import org.wildfly.common.function.ExceptionSupplier;
import org.wildfly.extension.datasources.agroal.logging.AgroalLogger;
import org.wildfly.extension.datasources.agroal.logging.LoggingDataSourceListener;
Expand Down Expand Up @@ -64,6 +63,7 @@
import java.security.AccessController;
import java.sql.Driver;
import java.sql.SQLException;
import java.util.function.Consumer;
import java.util.function.Supplier;

/**
Expand All @@ -75,6 +75,7 @@ public class DataSourceService implements Service<AgroalDataSource>, Supplier<Ag

private static final AuthenticationContextConfigurationClient AUTH_CONFIG_CLIENT = AccessController.doPrivileged(AuthenticationContextConfigurationClient.ACTION);

private final Consumer<AgroalDataSource> consumer;
private final String dataSourceName;
private final String jndiName;
private final boolean jta;
Expand All @@ -84,12 +85,22 @@ public class DataSourceService implements Service<AgroalDataSource>, Supplier<Ag
private final AgroalDataSourceConfigurationSupplier dataSourceConfiguration;
private AgroalDataSource agroalDataSource;

private InjectedValue<Class> driverInjector = new InjectedValue<>();
private InjectedValue<AuthenticationContext> authenticationContextInjector = new InjectedValue<>();
private InjectedValue<ExceptionSupplier<CredentialSource, Exception>> credentialSourceSupplierInjector = new InjectedValue<>();
private InjectedValue<TransactionSynchronizationRegistry> transactionSynchronizationRegistryInjector = new InjectedValue<>();

public DataSourceService(String dataSourceName, String jndiName, boolean jta, boolean connectable, boolean xa, AgroalDataSourceConfigurationSupplier dataSourceConfiguration) {
private final Supplier<Class> driverSupplier;
private final Supplier<AuthenticationContext> authenticationContextSupplier;
private final Supplier<ExceptionSupplier<CredentialSource, Exception>> credentialSourceSupplier;
private final Supplier<TransactionSynchronizationRegistry> transactionSynchronizationRegistrySupplier;

public DataSourceService(final Consumer<AgroalDataSource> consumer,
final Supplier<Class> driverSupplier,
final Supplier<AuthenticationContext> authenticationContextSupplier,
final Supplier<ExceptionSupplier<CredentialSource, Exception>> credentialSourceSupplier,
final Supplier<TransactionSynchronizationRegistry> transactionSynchronizationRegistrySupplier,
String dataSourceName, String jndiName, boolean jta, boolean connectable, boolean xa, AgroalDataSourceConfigurationSupplier dataSourceConfiguration) {
this.consumer = consumer;
this.driverSupplier = driverSupplier;
this.authenticationContextSupplier = authenticationContextSupplier;
this.credentialSourceSupplier = credentialSourceSupplier;
this.transactionSynchronizationRegistrySupplier = transactionSynchronizationRegistrySupplier;
this.dataSourceName = dataSourceName;
this.jndiName = jndiName;
this.jta = jta;
Expand All @@ -100,7 +111,7 @@ public DataSourceService(String dataSourceName, String jndiName, boolean jta, bo

@Override
public void start(StartContext context) throws StartException {
Class<?> providerClass = driverInjector.getOptionalValue();
Class<?> providerClass = driverSupplier != null ? driverSupplier.get() : null;
if (xa) {
if (!XADataSource.class.isAssignableFrom(providerClass)) {
throw AgroalLogger.SERVICE_LOGGER.invalidXAConnectionProvider();
Expand All @@ -115,7 +126,7 @@ public void start(StartContext context) throws StartException {

if (jta || xa) {
TransactionManager transactionManager = ContextTransactionManager.getInstance();
TransactionSynchronizationRegistry transactionSynchronizationRegistry = transactionSynchronizationRegistryInjector.getValue();
TransactionSynchronizationRegistry transactionSynchronizationRegistry = transactionSynchronizationRegistrySupplier != null ? transactionSynchronizationRegistrySupplier.get() : null;

if (transactionManager == null || transactionSynchronizationRegistry == null) {
throw AgroalLogger.SERVICE_LOGGER.missingTransactionManager();
Expand All @@ -124,7 +135,7 @@ public void start(StartContext context) throws StartException {
dataSourceConfiguration.connectionPoolConfiguration().transactionIntegration(txIntegration);
}

AuthenticationContext authenticationContext = authenticationContextInjector.getOptionalValue();
AuthenticationContext authenticationContext = authenticationContextSupplier != null ? authenticationContextSupplier.get() : null;

if (authenticationContext != null) {
try {
Expand Down Expand Up @@ -162,7 +173,7 @@ public void start(StartContext context) throws StartException {
}
}

ExceptionSupplier<CredentialSource, Exception> credentialSourceExceptionExceptionSupplier = credentialSourceSupplierInjector.getOptionalValue();
ExceptionSupplier<CredentialSource, Exception> credentialSourceExceptionExceptionSupplier = credentialSourceSupplier != null ? credentialSourceSupplier.get() : null;

if (credentialSourceExceptionExceptionSupplier != null) {
try {
Expand Down Expand Up @@ -196,10 +207,12 @@ public void start(StartContext context) throws StartException {
throw AgroalLogger.SERVICE_LOGGER.datasourceStartException(e, dataSourceName);
}
}
consumer.accept(agroalDataSource);
}

@Override
public void stop(StopContext context) {
consumer.accept(null);
agroalDataSource.close();
if (xa) {
AgroalLogger.SERVICE_LOGGER.stoppedXADataSource(dataSourceName);
Expand All @@ -218,21 +231,4 @@ public AgroalDataSource get() {
return agroalDataSource;
}

// --- //

public InjectedValue<Class> getDriverInjector() {
return this.driverInjector;
}

public InjectedValue<AuthenticationContext> getAuthenticationContextInjector() {
return this.authenticationContextInjector;
}

public InjectedValue<ExceptionSupplier<CredentialSource, Exception>> getCredentialSourceSupplierInjector() {
return credentialSourceSupplierInjector;
}

InjectedValue<TransactionSynchronizationRegistry> getTransactionSynchronizationRegistryInjector() {
return transactionSynchronizationRegistryInjector;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@

import static org.jboss.as.controller.descriptions.ModelDescriptionConstants.OP_ADDR;

import io.agroal.api.AgroalDataSource;
import jakarta.transaction.TransactionSynchronizationRegistry;

import io.agroal.api.configuration.supplier.AgroalConnectionFactoryConfigurationSupplier;
Expand All @@ -37,6 +38,12 @@
import org.jboss.as.controller.PathAddress;
import org.jboss.dmr.ModelNode;
import org.jboss.msc.service.ServiceName;
import org.wildfly.common.function.ExceptionSupplier;
import org.wildfly.security.auth.client.AuthenticationContext;
import org.wildfly.security.credential.source.CredentialSource;

import java.util.function.Consumer;
import java.util.function.Supplier;

/**
* Operations for adding and removing an xa-datasource resource to the model
Expand Down Expand Up @@ -80,17 +87,15 @@ protected void performRuntime(OperationContext context, ModelNode operation, Mod

String jndiName = AbstractDataSourceDefinition.JNDI_NAME_ATTRIBUTE.resolveModelAttribute(context, model).asString();
String driverName = AbstractDataSourceDefinition.DRIVER_ATTRIBUTE.resolveModelAttribute(context, factoryModel).asString();

DataSourceService dataSourceService = new DataSourceService(datasourceName, jndiName, false, false, true, dataSourceConfiguration);

CapabilityServiceBuilder serviceBuilder = context.getCapabilityServiceTarget().addCapability(AbstractDataSourceDefinition.DATA_SOURCE_CAPABILITY.fromBaseCapability(datasourceName))
.setInstance(dataSourceService)
.addCapabilityRequirement(DriverDefinition.AGROAL_DRIVER_CAPABILITY.getDynamicName(driverName), Class.class, dataSourceService.getDriverInjector());
CapabilityServiceBuilder serviceBuilder = context.getCapabilityServiceTarget().addCapability(AbstractDataSourceDefinition.DATA_SOURCE_CAPABILITY.fromBaseCapability(datasourceName));
final Consumer<AgroalDataSource> consumer = serviceBuilder.provides(AbstractDataSourceDefinition.DATA_SOURCE_CAPABILITY.fromBaseCapability(datasourceName));
final Supplier<Class> driverSupplier = serviceBuilder.requiresCapability(DriverDefinition.AGROAL_DRIVER_CAPABILITY.getDynamicName(driverName), Class.class);
final Supplier<AuthenticationContext> authenticationContextSupplier = AbstractDataSourceOperations.setupAuthenticationContext(context, factoryModel, serviceBuilder);
final Supplier<ExceptionSupplier<CredentialSource, Exception>> credentialSourceSupplier = AbstractDataSourceOperations.setupCredentialReference(context, factoryModel, serviceBuilder);
// TODO add a Stage.MODEL requirement
serviceBuilder.addCapabilityRequirement("org.wildfly.transactions.transaction-synchronization-registry", TransactionSynchronizationRegistry.class, dataSourceService.getTransactionSynchronizationRegistryInjector());

AbstractDataSourceOperations.setupElytronSecurity(context, factoryModel, dataSourceService, serviceBuilder);

final Supplier<TransactionSynchronizationRegistry> txnRegistrySupplier = serviceBuilder.requiresCapability("org.wildfly.transactions.transaction-synchronization-registry", TransactionSynchronizationRegistry.class);
DataSourceService dataSourceService = new DataSourceService(consumer, driverSupplier, authenticationContextSupplier, credentialSourceSupplier, txnRegistrySupplier, datasourceName, jndiName, false, false, true, dataSourceConfiguration);
serviceBuilder.setInstance(dataSourceService);
serviceBuilder.install();
}
}
Expand Down

0 comments on commit 2feaab5

Please sign in to comment.