Updates for running tests with managed identity (#2416)

Signed-off-by: Jeff Wasty <[email protected]>

# Conflicts:
#	src/test/java/com/microsoft/sqlserver/jdbc/SQLServerConnectionTest.java
#	src/test/java/com/microsoft/sqlserver/jdbc/TestResource.java
#	src/test/java/com/microsoft/sqlserver/jdbc/TestUtils.java
#	src/test/java/com/microsoft/sqlserver/jdbc/connection/TimeoutTest.java
#	src/test/java/com/microsoft/sqlserver/jdbc/databasemetadata/DatabaseMetaDataTest.java
#	src/test/java/com/microsoft/sqlserver/jdbc/unit/statement/BatchExecutionTest.java

Author: lilgreenbird

src/test/java/com/microsoft/sqlserver/jdbc/AlwaysEncrypted/AESetup.java

@@ -328,7 +328,7 @@ protected static void createTable(String tableName, String cekName, String table
             TestUtils.dropTableIfExists(tableName, stmt);
             sql = String.format(createSql, tableName, sql);
             stmt.execute(sql);
-            stmt.execute("DBCC FREEPROCCACHE");
+            TestUtils.freeProcCache(stmt);
         } catch (SQLException e) {
             fail(e.getMessage());
         }
@@ -362,7 +362,7 @@ protected static void createPrecisionTable(String tableName, String table[][], S
             }
             sql = String.format(createSql, tableName, sql);
             stmt.execute(sql);
-            stmt.execute("DBCC FREEPROCCACHE");
+            TestUtils.freeProcCache(stmt);
         } catch (SQLException e) {
             fail(e.getMessage());
         }
@@ -390,7 +390,7 @@ protected static void createScaleTable(String tableName, String table[][], Strin
 
             sql = String.format(createSql, tableName, sql);
             stmt.execute(sql);
-            stmt.execute("DBCC FREEPROCCACHE");
+            TestUtils.freeProcCache(stmt);
         } catch (SQLException e) {
             fail(e.getMessage());
         }

src/test/java/com/microsoft/sqlserver/jdbc/AlwaysEncrypted/CallableStatementTest.java

@@ -2200,7 +2200,7 @@ protected static void createDateTableCallableStatement(String cekName) throws SQ
                 SQLServerStatement stmt = (SQLServerStatement) con.createStatement()) {
             TestUtils.dropTableIfExists(DATE_TABLE_AE, stmt);
             stmt.execute(sql);
-            stmt.execute("DBCC FREEPROCCACHE");
+            TestUtils.freeProcCache(stmt);
         } catch (SQLException e) {
             fail(e.getMessage());
         }

src/test/java/com/microsoft/sqlserver/jdbc/SQLServerConnectionTest.java

@@ -43,6 +43,8 @@ protected Object[][] getContents() {
             {"R_lengthTruncated", " The inserted length is truncated or not correct!"},
             {"R_timeValueTruncated", " The time value is truncated or not correct!"},
             {"R_invalidErrorMessage", "Invalid Error Message: "},
+            {"R_kerberosNativeGSSFailure",
+                    "No valid credentials provided (Mechanism level: Failed to find any Kerberos tgt)"},
             {"R_expectedFailPassed", "Expected failure did not fail"}, {"R_dataTypeNotFound", "Cannot find data type"},
             {"R_illegalCharWktPosition", "Illegal character in Well-Known text at position {0}."},
             {"R_illegalCharWkt", "Illegal Well-Known text. Please make sure Well-Known text is valid."},
@@ -54,14 +56,23 @@ protected Object[][] getContents() {
             {"R_createDropAlterTableFailed", "Create/drop/alter table with preparedStatement failed!"},
             {"R_grantFailed", "grant table with preparedStatement failed!"},
             {"R_connectionIsClosed", "The connection is closed."},
+            {"R_noNamedAndIndexedParameters",
+                    "Cannot specify both named and indexed parameters when 'useFlexibleCallableStatements=false'"},
+            {"R_unknownOutputParameter",
+                    "Cannot acquire output parameter value by name. No parameter index was associated with the output parameter name. If acquiring output parameter by name, verify that the output parameter was initially registered by name."},
             {"R_ConnectionURLNull", "The connection URL is null."},
             {"R_connectionIsNotClosed", "The connection is not closed."},
             {"R_invalidExceptionMessage", "Invalid exception message"},
+            {"R_invalidClientSecret", "AADSTS7000215: Invalid client secret provided"},
+            {"R_invalidCertFields",
+                    "Error reading certificate, please verify the location of the certificate.signed fields invalid"},
+            {"R_invalidAADAuth", "Failed to authenticate the user {0} in Active Directory (Authentication={1})"},
             {"R_failedValidate", "failed to validate values in $0} "}, {"R_tableNotDropped", "table not dropped. "},
             {"R_connectionReset", "Connection reset"}, {"R_unknownException", "Unknown exception"},
             {"R_deadConnection", "Dead connection should be invalid"},
             {"R_wrongExceptionMessage", "Wrong exception message"}, {"R_wrongSqlState", "Wrong sql state"},
             {"R_parameterNotDefined", "Parameter {0} was not defined"},
+            {"R_notValidParameterForProcedure", "{0} is not a parameter for procedure {1}."},
             {"R_unexpectedExceptionContent", "Unexpected content in exception message"},
             {"R_connectionClosed", "The connection has been closed"},
             {"R_conversionFailed", "Conversion failed when converting {0} to {1} data type"},
@@ -201,5 +212,14 @@ protected Object[][] getContents() {
             {"R_objectNullOrEmpty", "The {0} is null or empty."},
             {"R_cekDecryptionFailed", "Failed to decrypt a column encryption key using key store provider: {0}."},
             {"R_connectTimedOut", "connect timed out"},
-            {"R_sessionKilled", "Cannot continue the execution because the session is in the kill state"}};
+            {"R_sharedTimerStopOnNoRef", "SharedTimer should be stopped after all references are removed."},
+            {"R_sessionKilled", "Cannot continue the execution because the session is in the kill state"},
+            {"R_failedFedauth", "Failed to acquire fedauth token: "},
+            {"R_noLoginModulesConfiguredForJdbcDriver",
+                    "javax.security.auth.login.LoginException (No LoginModules configured for SQLJDBCDriver)"},
+            {"R_unexpectedThreadCount", "Thread count is higher than expected."},
+            {"R_expectedClassDoesNotMatchActualClass",
+                    "Expected column class {0} does not match actual column class {1} for column {2}."},
+            {"R_loginFailedMI", "Login failed for user '<token-identified principal>'"},
+            {"R_MInotAvailable", "Managed Identity authentication is not available"},};
 }

src/test/java/com/microsoft/sqlserver/jdbc/TestResource.java

@@ -13,6 +13,9 @@
 import java.io.File;
 import java.io.FileInputStream;
 import java.io.FileOutputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.lang.reflect.Field;
 import java.net.URI;
 import java.security.KeyStore;
 import java.security.cert.CertificateFactory;
@@ -26,11 +29,27 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Calendar;
+import java.util.Date;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Locale;
+import java.util.Properties;
 import java.util.ResourceBundle;
-
+import java.util.Set;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+
+import org.junit.Assert;
+
+import com.microsoft.aad.msal4j.ClientCredentialFactory;
+import com.microsoft.aad.msal4j.ClientCredentialParameters;
+import com.microsoft.aad.msal4j.ConfidentialClientApplication;
+import com.microsoft.aad.msal4j.IAuthenticationResult;
+import com.microsoft.aad.msal4j.IClientCredential;
 import com.microsoft.sqlserver.testframework.AbstractSQLGenerator;
+import com.microsoft.sqlserver.testframework.Constants;
 import com.microsoft.sqlserver.testframework.PrepUtil;
 import com.microsoft.sqlserver.testframework.sqlType.SqlBigInt;
 import com.microsoft.sqlserver.testframework.sqlType.SqlBinary;
@@ -75,6 +94,97 @@ public final class TestUtils {
     static final int ENGINE_EDITION_FOR_SQL_AZURE_DW = 6;
     static final int ENGINE_EDITION_FOR_SQL_AZURE_MI = 8;
 
+    public static final int TEST_TOKEN_EXPIRY_SECONDS = 120; // token expiry time in secs
+
+    public static String ACCESS_TOKEN_CALLBACK = null;
+
+    static String applicationKey;
+    static String applicationClientID;
+
+    static {
+        try (InputStream input = new FileInputStream(Constants.CONFIG_PROPERTIES_FILE)) {
+            Properties configProperties = new Properties();
+            configProperties.load(input);
+            applicationKey = configProperties.getProperty("applicationKey");
+            applicationClientID = configProperties.getProperty("applicationClientID");
+        } catch (IOException e) {
+            // No config file found
+        }
+    }
+
+    public static boolean expireTokenToggle = false;
+
+    public static final SQLServerAccessTokenCallback accessTokenCallback = new SQLServerAccessTokenCallback() {
+        @Override
+        public SqlAuthenticationToken getAccessToken(String spn, String stsurl) {
+            String scope = spn + "/.default";
+            Set<String> scopes = new HashSet<>();
+            scopes.add(scope);
+
+            try {
+                ExecutorService executorService = Executors.newSingleThreadExecutor();
+                IClientCredential credential = ClientCredentialFactory.createFromSecret(applicationKey);
+                ConfidentialClientApplication clientApplication = ConfidentialClientApplication
+                        .builder(applicationClientID, credential).executorService(executorService).authority(stsurl)
+                        .build();
+                CompletableFuture<IAuthenticationResult> future = clientApplication
+                        .acquireToken(ClientCredentialParameters.builder(scopes).build());
+
+                IAuthenticationResult authenticationResult = future.get();
+                String accessToken = authenticationResult.accessToken();
+                long expiresOn = authenticationResult.expiresOnDate().getTime();
+
+                ACCESS_TOKEN_CALLBACK = accessToken;
+
+                if (expireTokenToggle) {
+                    Date now = new Date();
+                    long minutesToExpireWithin = TEST_TOKEN_EXPIRY_SECONDS * 1000; // Expire within 2 minutes
+                    return new SqlAuthenticationToken(accessToken, now.getTime() + minutesToExpireWithin);
+                } else {
+                    return new SqlAuthenticationToken(accessToken, expiresOn);
+                }
+            } catch (Exception e) {
+                fail(TestResource.getResource("R_unexpectedException") + e.getMessage());
+            }
+            return null;
+        }
+    };
+
+    public static void setAccessTokenExpiry(Object con, String accessToken) {
+        Field fedAuthTokenField;
+        try {
+            fedAuthTokenField = SQLServerConnection.class.getDeclaredField("fedAuthToken");
+            fedAuthTokenField.setAccessible(true);
+
+            Date newExpiry = new Date(
+                    System.currentTimeMillis() + TimeUnit.SECONDS.toMillis(TEST_TOKEN_EXPIRY_SECONDS));
+            SqlAuthenticationToken newFedAuthToken = new SqlAuthenticationToken(accessToken, newExpiry);
+            fedAuthTokenField.set(con, newFedAuthToken);
+        } catch (NoSuchFieldException | SecurityException | IllegalArgumentException | IllegalAccessException e) {
+            Assert.fail("Failed to set token expiry: " + e.getMessage());
+        }
+    }
+
+    public static void setAccessTokenExpiry(Object con) {
+        Field fedAuthTokenField;
+        Field wrappedConnection;
+        try {
+            fedAuthTokenField = SQLServerConnection.class.getDeclaredField("fedAuthToken");
+            fedAuthTokenField.setAccessible(true);
+
+            wrappedConnection = SQLServerConnectionPoolProxy.class.getDeclaredField("wrappedConnection");
+            wrappedConnection.setAccessible(true);
+            Object wrappedConnectionObj = wrappedConnection.get(con);
+
+            Date newExpiry = new Date(
+                    System.currentTimeMillis() + TimeUnit.SECONDS.toMillis(TEST_TOKEN_EXPIRY_SECONDS));
+            SqlAuthenticationToken newFedAuthToken = new SqlAuthenticationToken(ACCESS_TOKEN_CALLBACK, newExpiry);
+            fedAuthTokenField.set(wrappedConnectionObj, newFedAuthToken);
+        } catch (NoSuchFieldException | SecurityException | IllegalArgumentException | IllegalAccessException e) {
+            Assert.fail("Failed to set token expiry: " + e.getMessage());
+        }
+    }
+
     private TestUtils() {}
 
     /**
@@ -127,6 +237,15 @@ public static boolean isAzureMI(Connection con) {
         return ((SQLServerConnection) con).isAzureMI();
     }
 
+    /**
+     * Checks if connection is established to Azure Synapse OnDemand server
+     * 
+     */
+    public static boolean isAzureSynapseOnDemand(Connection con) {
+        isAzure(con);
+        return ((SQLServerConnection) con).isAzureSynapseOnDemandEndpoint();
+    }
+
     /**
      * Checks if connection is established to server that supports AEv2.
      * 
@@ -300,6 +419,18 @@ public static void dropTableIfExists(String tableName, java.sql.Statement stmt)
         dropObjectIfExists(tableName, "U", stmt);
     }
 
+    public static void dropTableWithSchemaIfExists(String tableNameWithSchema,
+            java.sql.Statement stmt) throws SQLException {
+        stmt.execute(
+                "IF OBJECT_ID('" + tableNameWithSchema + "', 'U') IS NOT NULL DROP TABLE " + tableNameWithSchema + ";");
+    }
+
+    public static void dropProcedureWithSchemaIfExists(String procedureWithSchema,
+            java.sql.Statement stmt) throws SQLException {
+        stmt.execute("IF EXISTS (SELECT * FROM sys.objects WHERE object_id = OBJECT_ID(N'" + procedureWithSchema
+                + "') AND type in (N'P', N'PC')) DROP PROCEDURE " + procedureWithSchema + ";");
+    }
+
     /**
      * Deletes the contents of a table.
      * 
@@ -376,7 +507,8 @@ public static void dropTypeIfExists(String typeName, java.sql.Statement stmt) th
      * @throws SQLException
      */
     public static void dropUserDefinedTypeIfExists(String typeName, Statement stmt) throws SQLException {
-        stmt.executeUpdate("IF EXISTS (select * from sys.types where name = '" + escapeSingleQuotes(typeName) + "') DROP TYPE " + typeName);
+        stmt.executeUpdate("IF EXISTS (select * from sys.types where name = '" + escapeSingleQuotes(typeName)
+                + "') DROP TYPE " + typeName);
     }
 
     /**
@@ -403,7 +535,31 @@ public static void dropDatabaseIfExists(String databaseName, String connectionSt
      */
     public static void dropSchemaIfExists(String schemaName, Statement stmt) throws SQLException {
         stmt.execute("if EXISTS (SELECT * FROM sys.schemas where name = '" + escapeSingleQuotes(schemaName)
-                + "') drop schema " + AbstractSQLGenerator.escapeIdentifier(schemaName));
+                + "') DROP SCHEMA" + AbstractSQLGenerator.escapeIdentifier(schemaName));
+    }
+
+    /**
+     * mimic "DROP USER..."
+     * 
+     * @param userName
+     * @param stmt
+     * @throws SQLException
+     */
+    public static void dropUserIfExists(String userName, Statement stmt) throws SQLException {
+        stmt.execute("IF EXISTS (SELECT * FROM sys.sysusers where name = '" + escapeSingleQuotes(userName)
+                + "') DROP USER " + AbstractSQLGenerator.escapeIdentifier(userName));
+    }
+
+    /**
+     * mimic "DROP LOGIN..."
+     * 
+     * @param userName
+     * @param stmt
+     * @throws SQLException
+     */
+    public static void dropLoginIfExists(String userName, Statement stmt) throws SQLException {
+        stmt.execute("IF EXISTS (SELECT * FROM sys.sysusers where name = '" + escapeSingleQuotes(userName)
+                + "') DROP LOGIN " + AbstractSQLGenerator.escapeIdentifier(userName));
     }
 
     /**
@@ -964,4 +1120,27 @@ private static java.security.cert.Certificate getCertificate(String certname) th
             return cf.generateCertificate(is);
         }
     }
+
+    public static String getConnectionID(
+            SQLServerPooledConnection pc) throws ClassNotFoundException, NoSuchFieldException, IllegalAccessException {
+        Class<?> pooledConnection = Class.forName("com.microsoft.sqlserver.jdbc.SQLServerPooledConnection");
+        Class<?> connection = Class.forName("com.microsoft.sqlserver.jdbc.SQLServerConnection");
+
+        Field physicalConnection = pooledConnection.getDeclaredField("physicalConnection");
+        Field traceID = connection.getDeclaredField("traceID");
+
+        physicalConnection.setAccessible(true);
+        traceID.setAccessible(true);
+
+        SQLServerConnection conn = (SQLServerConnection) physicalConnection.get(pc);
+        return (String) traceID.get(conn);
+    }
+
+    public static void freeProcCache(Statement stmt) {
+        try {
+            stmt.execute("DBCC FREEPROCCACHE");
+        } catch (Exception e) {
+            // ignore error - some tests fails due to permission issues from managed identity, this does not seem to affect tests
+        }
+    }
 }

src/test/java/com/microsoft/sqlserver/jdbc/TestUtils.java

@@ -154,6 +154,10 @@ public void testConnectionPoolClose() throws SQLException {
 
     @Test
     public void testConnectionPoolClientConnectionId() throws SQLException {
+        String auth = TestUtils.getProperty(connectionString, "authentication");
+        org.junit.Assume.assumeTrue(auth != null
+                && (auth.equalsIgnoreCase("SqlPassword") || auth.equalsIgnoreCase("ActiveDirectoryPassword")));
+
         SQLServerXADataSource ds = new SQLServerXADataSource();
         ds.setURL(connectionString);
         PooledConnection pc = null;