exceptionfactory commented on code in PR #11722:
URL: https://github.com/apache/nifi/pull/11722#discussion_r4136547888


##########
nifi-extension-bundles/nifi-extension-utils/nifi-dbcp-base/src/test/java/org/apache/nifi/dbcp/AbstractDBCPConnectionPoolTest.java:
##########
@@ -132,16 +171,172 @@ void testVerifyUsesDatabasePasswordProvider() throws 
SQLException {
         
when(configurationContext.getProperty(eq(DBCPProperties.PASSWORD_SOURCE))).thenReturn(propertyValue(DBCPProperties.PASSWORD_SOURCE,
                 DBCPProperties.PasswordSource.PASSWORD_PROVIDER.getValue()));
         when(connection.isValid(eq(TIMEOUT))).thenReturn(true);
-        
when(databasePasswordProvider.getPassword(any())).thenAnswer(invocation -> 
"token".toCharArray());
-
-        final ArgumentCaptor<Properties> propertiesCaptor = 
ArgumentCaptor.forClass(Properties.class);
-        when(driver.connect(any(), 
propertiesCaptor.capture())).thenReturn(connection);
+        
when(databasePasswordProvider.getDatabaseCredentialPlacement()).thenAnswer(invocation
 -> credentialPlacement.get());
+        
when(databasePasswordProvider.getPassword(any())).thenAnswer(invocation -> {
+            final char[] credential = 
credentialValues[issuedCredentialIndex.getAndIncrement()].toCharArray();
+            issuedCredentials.add(credential);
+            return credential;
+        });
+        when(driver.connect(any(), any())).thenAnswer(invocation -> {
+            final Properties properties = invocation.getArgument(1);
+            final Properties snapshot = new Properties();
+            snapshot.putAll(properties);
+
+            if ("VERIFY".equals(connectionMode.get())) {
+                return connection;
+            }
+
+            if ("PASSWORD".equals(connectionMode.get())) {
+                passwordSnapshots.add(snapshot);
+                passwordLiveReferences.add(properties);
+                passwordConnectEntered.countDown();
+                assertTrue(releasePasswordConnect.await(5, TimeUnit.SECONDS));
+                return connection;
+            }
+
+            if (accessTokenAttempt.getAndIncrement() == 0) {
+                accessTokenSuccessSnapshot.set(snapshot);
+                accessTokenSuccessLiveReference.set(properties);
+                return connection;
+            }
+
+            accessTokenFailureSnapshot.set(snapshot);
+            accessTokenFailureLiveReference.set(properties);
+            throw new SQLException("driver failure");
+        });
 
         connectionPool.verify(configurationContext, componentLog, 
Collections.emptyMap());
+        issuedCredentialIndex.set(0);
+        issuedCredentials.clear();
+        connectionMode.set("PASSWORD");
 
-        final List<Properties> capturedProperties = 
propertiesCaptor.getAllValues();
-        assertTrue(capturedProperties.stream().anyMatch(props -> 
"token".equals(props.getProperty("password"))));
         verify(databasePasswordProvider, atLeastOnce()).getPassword(any());
+        verify(databasePasswordProvider, 
atLeastOnce()).getDatabaseCredentialPlacement();
+
+        passwordDataSource.setDriver(driver);
+        passwordDataSource.setUrl("jdbc:postgresql://example");
+        passwordDataSource.setUsername("dbuser");
+        passwordDataSource.setPassword("configured-password");
+        passwordDataSource.setMaxTotal(MAX_TOTAL);
+        passwordDataSource.addConnectionProperty("ssl", "true");
+        
passwordDataSource.setDatabasePasswordProvider(databasePasswordProvider, 
DatabasePasswordRequestContext.builder()
+                .jdbcUrl("jdbc:postgresql://example")
+                .driverClassName("org.postgresql.Driver")
+                .databaseUser("dbuser")
+                .connectionProperties(Map.of("ssl", "true"))
+                .build());
+
+        final ConnectionFactory passwordConnectionFactory = 
passwordDataSource.callCreateConnectionFactory();
+        final ExecutorService executorService = 
Executors.newFixedThreadPool(2);
+        try {
+            final Future<Connection> firstConnection = 
executorService.submit(passwordConnectionFactory::createConnection);
+            final Future<Connection> secondConnection = 
executorService.submit(passwordConnectionFactory::createConnection);
+
+            assertTrue(passwordConnectEntered.await(5, TimeUnit.SECONDS));
+            releasePasswordConnect.countDown();
+
+            firstConnection.get(5, TimeUnit.SECONDS).close();
+            secondConnection.get(5, TimeUnit.SECONDS).close();
+        } catch (final Exception e) {
+            throw new AssertionError(e);
+        } finally {
+            executorService.shutdownNow();
+            passwordDataSource.close();
+        }
+
+        assertEquals(2, passwordSnapshots.size());
+        assertNotSame(passwordLiveReferences.get(0), 
passwordLiveReferences.get(1));
+        assertTrue(passwordSnapshots.stream().allMatch(properties -> 
"true".equals(properties.getProperty("ssl"))));
+        assertTrue(passwordSnapshots.stream().anyMatch(properties -> 
"password-token-1".equals(properties.getProperty("password"))));
+        assertTrue(passwordSnapshots.stream().anyMatch(properties -> 
"password-token-2".equals(properties.getProperty("password"))));
+        assertTrue(passwordLiveReferences.stream().allMatch(properties -> 
properties.getProperty("password") == null));
+
+        credentialPlacement.set(DatabaseCredentialPlacement.ACCESS_TOKEN);
+        connectionMode.set("ACCESS");
+        accessTokenDataSource.setDriver(driver);
+        
accessTokenDataSource.setUrl("jdbc:sqlserver://example.database.windows.net:1433;databaseName=test");
+        accessTokenDataSource.setUsername("configured-user");
+        accessTokenDataSource.setPassword("configured-password");
+        accessTokenDataSource.addConnectionProperty("USER", "dynamic-user");
+        accessTokenDataSource.addConnectionProperty("UserName", 
"dynamic-user-name");
+        accessTokenDataSource.addConnectionProperty("PASSWORD", 
"dynamic-password");
+        accessTokenDataSource.addConnectionProperty("authentication", 
"ActiveDirectoryManagedIdentity");
+        accessTokenDataSource.addConnectionProperty("integratedSecurity", 
"true");
+        
accessTokenDataSource.setDatabasePasswordProvider(databasePasswordProvider, 
DatabasePasswordRequestContext.builder()
+                
.jdbcUrl("jdbc:sqlserver://example.database.windows.net:1433;databaseName=test")
+                
.driverClassName("com.microsoft.sqlserver.jdbc.SQLServerDriver")
+                .databaseUser("configured-user")
+                .connectionProperties(Map.of(
+                        "USER", "dynamic-user",
+                        "UserName", "dynamic-user-name",
+                        "PASSWORD", "dynamic-password",
+                        "authentication", "ActiveDirectoryManagedIdentity",
+                        "integratedSecurity", "true"

Review Comment:
   A number of these properties and values appear to be repeated across test 
methods, so it would be helpful to move them into static final variables



##########
nifi-extension-bundles/nifi-extension-utils/nifi-dbcp-base/src/main/java/org/apache/nifi/dbcp/ProviderAwareBasicDataSource.java:
##########
@@ -20,17 +20,24 @@
 import org.apache.commons.dbcp2.ConnectionFactory;
 import org.apache.commons.dbcp2.Constants;
 import org.apache.commons.dbcp2.DriverConnectionFactory;
+import org.apache.nifi.dbcp.api.DatabaseCredentialPlacement;
 import org.apache.nifi.dbcp.api.DatabasePasswordProvider;
 import org.apache.nifi.dbcp.api.DatabasePasswordRequestContext;
 
 import java.sql.Connection;
 import java.sql.SQLException;
 import java.util.Arrays;
+import java.util.Locale;
+import java.util.Properties;
 
 /**
  * Extension of {@link BasicDataSource} that supports obtaining database 
passwords from a {@link DatabasePasswordProvider}.
  */
 public class ProviderAwareBasicDataSource extends BasicDataSource {
+    private static final String USER_PROPERTY = "user";
+    private static final String USER_NAME_PROPERTY = "username";
+    private static final String ACCESS_TOKEN_PROPERTY = "accessToken";

Review Comment:
   How generic is this `accessToken` property? It seems more specific to MS SQL 
Server. With the Placement enum, it looks like a path forward for future 
additional options, but do you expect this to apply to other types of databases?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]

Reply via email to