Repository: microsoft/mssql-jdbc
Pull request #2073 (2023-02-10 00:12:33)

Added access token callback class connection string property
Addresses https://github.com/microsoft/mssql-jdbc/issues/2058

Added new `accessTokenCallbackClass` connection string property. The callback can be set by specifying the fully qualified name of the callback implementing class in the connection string.


Patch:
diff --git a/build.gradle b/build.gradle
index 4a3be342..9b7652ad 100644
--- a/build.gradle
+++ b/build.gradle
@@ -128,7 +128,7 @@ repositories {
 
 dependencies {
     implementation 'org.osgi:org.osgi.core:6.0.0',
-			'org.osgi:org.osgi.compendium:5.0.0'
+			'org.osgi:org.osgi.service.jdbc:1.1.0'
 	compileOnly 'com.azure:azure-security-keyvault-keys:4.5.2',
 			'com.azure:azure-identity:1.7.0',
 			'org.antlr:antlr4-runtime:4.9.3',
diff --git a/pom.xml b/pom.xml
index a4420cfd..37540dac 100644
--- a/pom.xml
+++ b/pom.xml
@@ -66,7 +66,7 @@
 		<azure-security-keyvault-keys.version>4.5.3</azure-security-keyvault-keys.version>
 		<azure-identity.version>1.7.0</azure-identity.version>
 		<msal.version>1.13.3</msal.version>
-		<org.osgi.compendium.version>5.0.0</org.osgi.compendium.version>
+		<osgi.jdbc.version>1.1.0</osgi.jdbc.version>
 		<antlr-runtime.version>4.9.3</antlr-runtime.version>
 		<com.google.code.gson.version>2.9.0</com.google.code.gson.version>
 		<bcprov-jdk15on.version>1.70</bcprov-jdk15on.version>
@@ -150,8 +150,8 @@
 		</dependency>
 		<dependency>
 			<groupId>org.osgi</groupId>
-			<artifactId>org.osgi.compendium</artifactId>
-			<version>${org.osgi.compendium.version}</version>
+			<artifactId>org.osgi.service.jdbc</artifactId>
+			<version>${osgi.jdbc.version}</version>
 			<scope>provided</scope>
 		</dependency>
 
@@ -495,7 +495,7 @@
 							com.microsoft.sqlserver.jdbc.dataclassification,
 							microsoft.sql
 						</_exportcontents>
-						<Import-Package>!microsoft.sql,*</Import-Package>
+						<Import-Package>!microsoft.sql,jdk.net;resolution:=optional,*</Import-Package>
 						<Bundle-Activator>com.microsoft.sqlserver.jdbc.osgi.Activator</Bundle-Activator>
 					</instructions>
 				</configuration>
diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/ISQLServerConnection.java b/src/main/java/com/microsoft/sqlserver/jdbc/ISQLServerConnection.java
index ee0a5a8e..88fa73a3 100644
--- a/src/main/java/com/microsoft/sqlserver/jdbc/ISQLServerConnection.java
+++ b/src/main/java/com/microsoft/sqlserver/jdbc/ISQLServerConnection.java
@@ -240,7 +240,7 @@ public interface ISQLServerConnection extends java.sql.Connection {
      * Valid values are: datetime, datetime2 or datetimeoffset.
      * 
      * @param datetimeParameterTypeValue
-     *        The datatype to use when encoding Java dates into SQL Server. Valid values are: 
+     *        The datatype to use when encoding Java dates into SQL Server. Valid values are:
      *        datetime, datetime2 or datetimeoffset.
      * 
      * @throws SQLServerException
@@ -451,4 +451,19 @@ public interface ISQLServerConnection extends java.sql.Connection {
      */
     @Deprecated(since = "12.1.0", forRemoval = true)
     void setMsiTokenCacheTtl(int timeToLive);
+
+    /**
+     * Returns the fully qualified class name of the implementing class for {@link SQLServerAccessTokenCallback}.
+     *
+     * @return accessTokenCallbackClass
+     */
+    String getAccessTokenCallbackClass();
+
+    /**
+     * Sets 'accessTokenCallbackClass' to the fully qualified class name
+     * of the implementing class for {@link SQLServerAccessTokenCallback}.
+     *
+     * @param accessTokenCallbackClass
+     */
+    void setAccessTokenCallbackClass(String accessTokenCallbackClass);
 }
diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/ISQLServerDataSource.java b/src/main/java/com/microsoft/sqlserver/jdbc/ISQLServerDataSource.java
index 4e6c8390..a435666e 100644
--- a/src/main/java/com/microsoft/sqlserver/jdbc/ISQLServerDataSource.java
+++ b/src/main/java/com/microsoft/sqlserver/jdbc/ISQLServerDataSource.java
@@ -347,7 +347,7 @@ public interface ISQLServerDataSource extends javax.sql.CommonDataSource {
      * Sets the SQL server datatype to use for Java datetime and timestamp values.
      * 
      * @param datetimeParameterType
-     *        The SQL datatype to use when encoding Java dates for SQL Server. Valid values are: 
+     *        The SQL datatype to use when encoding Java dates for SQL Server. Valid values are:
      *        datetime, datetime2 or datetimeoffset.
      */
     void setDatetimeParameterType(String datetimeParameterType);
@@ -1268,4 +1268,19 @@ public interface ISQLServerDataSource extends javax.sql.CommonDataSource {
      * @return Access token callback delegate.
      */
     SQLServerAccessTokenCallback getAccessTokenCallback();
+
+    /**
+     * Returns the fully qualified class name of the implementing class for {@link SQLServerAccessTokenCallback}.
+     *
+     * @return accessTokenCallbackClass
+     */
+    String getAccessTokenCallbackClass();
+
+    /**
+     * Sets 'accessTokenCallbackClass' to the fully qualified class name
+     * of the implementing class for {@link SQLServerAccessTokenCallback}.
+     *
+     * @param accessTokenCallbackClass
+     */
+    void setAccessTokenCallbackClass(String accessTokenCallbackClass);
 }
diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerAccessTokenCallback.java b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerAccessTokenCallback.java
index 2d3ba451..c56c6be4 100644
--- a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerAccessTokenCallback.java
+++ b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerAccessTokenCallback.java
@@ -17,5 +17,5 @@ public interface SQLServerAccessTokenCallback {
      *
      * @return Returns a {@link SqlAuthenticationToken}.
      */
-    SqlAuthenticationToken getAccessToken(String stsurl, String spn);
+    SqlAuthenticationToken getAccessToken(String spn, String stsurl);
 }
diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerCallableStatement.java b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerCallableStatement.java
index 7b332f0a..ee6be290 100644
--- a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerCallableStatement.java
+++ b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerCallableStatement.java
@@ -1362,10 +1362,10 @@ public class SQLServerCallableStatement extends SQLServerPreparedStatement imple
 
         }
 
-        // @RETURN_VALUE will always be in the parameterNames map, so parameterNamesSize will always be at least of size 1.
-        // If the server didn't return anything (eg. the param names for the sproc), user might not have access.
-        // So, parameterNamesSize must be of size 1.
-        if (null != parameterNames && parameterNames.size() == 1) {
+        // If the server didn't return anything (eg. the param names for the sp_sproc_columns), user might not
+        // have required permissions to view all the parameterNames. And, there's also the case depending on the permissions,
+        // @RETURN_VALUE may or may not be present. So, the parameterNames list might have an additional +1 parameter.
+        if (null != parameterNames && parameterNames.size() <= 1) {
             return map.computeIfAbsent(columnName, ifAbsent -> ai.incrementAndGet());
         }
 
diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerConnection.java b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerConnection.java
index e826ea0f..3ee35c2a 100644
--- a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerConnection.java
+++ b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerConnection.java
@@ -144,6 +144,15 @@ public class SQLServerConnection implements ISQLServerConnection, java.io.Serial
     /** fedAuth requested flag */
     private boolean federatedAuthenticationRequested = false;
 
+    /** Access token callback class name */
+    private String accessTokenCallbackClass = null;
+
+    /** Flag that determines whether the accessToken callback class is set **/
+    private boolean hasAccessTokenCallbackClass = false;
+
+    /** Flag that determines whether the accessToken callback was set **/
+    private SQLServerAccessTokenCallback accessTokenCallback = null;
+
     /**
      * Keep this distinct from _federatedAuthenticationRequested, since some fedauth library types may not need more
      * info
@@ -206,6 +215,9 @@ public class SQLServerConnection implements ISQLServerConnection, java.io.Serial
     /** flag indicating whether server is Azure MI */
     private Boolean isAzureMI = null;
 
+    /** flag indicating whether server supports transactions */
+    private Boolean supportsTransactions = null;
+
     /** shared timer */
     private SharedTimer sharedTimer;
 
@@ -536,12 +548,10 @@ public class SQLServerConnection implements ISQLServerConnection, java.io.Serial
                 default:
                     // If authenticationString not specified, check if access token callback was set.
                     // If access token callback is set, break.
-                    if (null != activeConnectionProperties
-                            .get(SQLServerDriverObjectProperty.ACCESS_TOKEN_CALLBACK.toString())) {
+                    if (null != accessTokenCallback || hasAccessTokenCallbackClass) {
                         this.authentication = SqlAuthentication.NOT_SPECIFIED;
                         break;
                     }
-                    assert (false);
                     MessageFormat form = new MessageFormat(
                             SQLServerException.getErrString("R_InvalidConnectionSetting"));
                     Object[] msgArgs = {"authentication", authenticationString};
@@ -2449,10 +2459,15 @@ public class SQLServerConnection implements ISQLServerConnection, java.io.Serial
                     ntlmAuthentication = true;
                 }
 
-                SQLServerAccessTokenCallback callback = (SQLServerAccessTokenCallback) activeConnectionProperties
+                accessTokenCallback = (SQLServerAccessTokenCallback) activeConnectionProperties
                         .get(SQLServerDriverObjectProperty.ACCESS_TOKEN_CALLBACK.toString());
 
-                if (null != callback && (!activeConnectionProperties
+                hasAccessTokenCallbackClass = null != activeConnectionProperties
+                        .get(SQLServerDriverStringProperty.ACCESS_TOKEN_CALLBACK_CLASS.toString())
+                        && !activeConnectionProperties
+                                .getProperty(SQLServerDriverStringProperty.ACCESS_TOKEN_CALLBACK_CLASS.toString())
+                                .isEmpty();
+                if ((null != accessTokenCallback || hasAccessTokenCallbackClass) && (!activeConnectionProperties
                         .getProperty(SQLServerDriverStringProperty.USER.toString()).isEmpty()
                         || !activeConnectionProperties.getProperty(SQLServerDriverStringProperty.PASSWORD.toString())
                                 .isEmpty())) {
@@ -2460,6 +2475,14 @@ public class SQLServerConnection implements ISQLServerConnection, java.io.Serial
                             SQLServerException.getErrString("R_AccessTokenCallbackWithUserPassword"), null);
                 }
 
+                sPropKey = SQLServerDriverStringProperty.ACCESS_TOKEN_CALLBACK_CLASS.toString();
+                sPropValue = activeConnectionProperties.getProperty(sPropKey);
+                if (null == sPropValue) {
+                    sPropValue = SQLServerDriverStringProperty.ACCESS_TOKEN_CALLBACK_CLASS.getDefaultValue();
+                    activeConnectionProperties.setProperty(sPropKey, sPropValue);
+                }
+                setAccessTokenCallbackClass(sPropValue);
+
                 sPropKey = SQLServerDriverStringProperty.AUTHENTICATION.toString();
                 sPropValue = activeConnectionProperties.getProperty(sPropKey);
                 if (null == sPropValue) {
@@ -3500,8 +3523,7 @@ public class SQLServerConnection implements ISQLServerConnection, java.io.Serial
     void prelogin(String serverName, int portNumber) throws SQLServerException {
         // Build a TDS Pre-Login packet to send to the server.
         if ((!authenticationString.equalsIgnoreCase(SqlAuthentication.NOT_SPECIFIED.toString()))
-                || (null != accessTokenInByte) || null != activeConnectionProperties
-                        .get(SQLServerDriverObjectProperty.ACCESS_TOKEN_CALLBACK.toString())) {
+                || (null != accessTokenInByte) || null != accessTokenCallback || hasAccessTokenCallbackClass) {
             fedAuthRequiredByUser = true;
         }
 
@@ -3883,8 +3905,8 @@ public class SQLServerConnection implements ISQLServerConnection, java.io.Serial
                     // Or AccessToken is not null, mean token based authentication is used.
                     if (((null != authenticationString)
                             && (!authenticationString.equalsIgnoreCase(SqlAuthentication.NOT_SPECIFIED.toString())))
-                            || (null != accessTokenInByte) || null != activeConnectionProperties
-                                    .get(SQLServerDriverObjectProperty.ACCESS_TOKEN_CALLBACK.toString())) {
+                            || (null != accessTokenInByte) || null != accessTokenCallback
+                            || hasAccessTokenCallbackClass) {
                         fedAuthRequiredPreLoginResponse = (preloginResponse[optionOffset] == 1);
                     }
                     break;
@@ -4827,8 +4849,7 @@ public class SQLServerConnection implements ISQLServerConnection, java.io.Serial
                             break;
                         default:
                             // If not specified, check if access token callback was set. If it is set, break.
-                            if (null != activeConnectionProperties
-                                    .get(SQLServerDriverObjectProperty.ACCESS_TOKEN_CALLBACK.toString())) {
+                            if (null != accessTokenCallback || hasAccessTokenCallbackClass) {
                                 workflow = TDS.ADALWORKFLOW_ACCESSTOKENCALLBACK;
                                 break;
                             }
@@ -5048,8 +5069,7 @@ public class SQLServerConnection implements ISQLServerConnection, java.io.Serial
                         || authenticationString
                                 .equalsIgnoreCase(SqlAuthentication.ACTIVE_DIRECTORY_INTERACTIVE.toString()))
                         && fedAuthRequiredPreLoginResponse)
-                || null != activeConnectionProperties
-                        .get(SQLServerDriverObjectProperty.ACCESS_TOKEN_CALLBACK.toString())) {
+                || null != accessTokenCallback || hasAccessTokenCallbackClass) {
             federatedAuthenticationInfoRequested = true;
             fedAuthFeatureExtensionData = new FederatedAuthenticationFeatureExtensionData(TDS.TDS_FEDAUTH_LIBRARY_ADAL,
                     authenticationString, fedAuthRequiredPreLoginResponse);
@@ -5609,11 +5629,22 @@ public class SQLServerConnection implements ISQLServerConnection, java.io.Serial
 
         attemptRefreshTokenLocked = true;
 
-        SQLServerAccessTokenCallback callback = (SQLServerAccessTokenCallback) activeConnectionProperties
-                .get(SQLServerDriverObjectProperty.ACCESS_TOKEN_CALLBACK.toString());
-
-        if (authenticationString.equals(SqlAuthentication.NOT_SPECIFIED.toString()) && null != callback) {
-            fedAuthToken = callback.getAccessToken(fedAuthInfo.spn, fedAuthInfo.stsurl);
+        if (authenticationString.equals(SqlAuthentication.NOT_SPECIFIED.toString()) && null != accessTokenCallbackClass
+                && !accessTokenCallbackClass.isEmpty()) {
+            try {
+                Object[] msgArgs = {"accessTokenCallbackClass",
+                        "com.microsoft.sqlserver.jdbc.SQLServerAccessTokenCallback"};
+                SQLServerAccessTokenCallback callbackInstance = Util.newInstance(SQLServerAccessTokenCallback.class,
+                        accessTokenCallbackClass, null, msgArgs);
+                fedAuthToken = callbackInstance.getAccessToken(fedAuthInfo.spn, fedAuthInfo.stsurl);
+            } catch (Exception e) {
+                MessageFormat form = new MessageFormat(
+                        SQLServerException.getErrString("R_InvalidAccessTokenCallbackClass"));
+                throw new SQLServerException(form.format(new Object[] {accessTokenCallbackClass}), e.getCause());
+            }
+        } else if (authenticationString.equals(SqlAuthentication.NOT_SPECIFIED.toString())
+                && null != accessTokenCallback) {
+            fedAuthToken = accessTokenCallback.getAccessToken(fedAuthInfo.spn, fedAuthInfo.stsurl);
         } else {
             fedAuthToken = getFedAuthToken(fedAuthInfo);
         }
@@ -7619,6 +7650,29 @@ public class SQLServerConnection implements ISQLServerConnection, java.io.Serial
         return 1 < getServerPreparedStatementDiscardThreshold();
     }
 
+    /**
+     * Returns the fully qualified class name of the implementing class for {@link SQLServerAccessTokenCallback}.
+     *
+     * @return accessTokenCallbackClass
+     */
+    public String getAccessTokenCallbackClass() {
+        if (null == this.accessTokenCallbackClass) {
+            return SQLServerDriverStringProperty.ACCESS_TOKEN_CALLBACK_CLASS.getDefaultValue();
+        }
+
+        return this.accessTokenCallbackClass;
+    };
+
+    /**
+     * Sets 'accessTokenCallbackClass' to the fully qualified class name
+     * of the implementing class for {@link SQLServerAccessTokenCallback}.
+     *
+     * @param accessTokenCallbackClass
+     */
+    public void setAccessTokenCallbackClass(String accessTokenCallbackClass) {
+        this.accessTokenCallbackClass = accessTokenCallbackClass;
+    };
+
     /**
      * Cleans up discarded prepared statement handles on the server using batched un-prepare actions if the batching
      * threshold has been reached.
@@ -7894,6 +7948,28 @@ public class SQLServerConnection implements ISQLServerConnection, java.io.Serial
         return isAzureMI;
     }
 
+    /**
+     * Checks if the connection established to server supports transactions.
+     *
+     * @return True if server supports transactions, otherwise false
+     */
+    boolean supportsTransactions() throws SQLServerException {
+        if (supportsTransactions != null) {
+            return supportsTransactions;
+        }
+
+        try {
+            this.connectionCommand("SELECT @@TRANCOUNT", "SQLServerConnection.supportsTransactions");
+        } catch (SQLServerException e) {
+            if (e.getMessage().equals(SQLServerException.getErrString("R_transactionsNotSupported"))) {
+                return supportsTransactions = false;
+            }
+            throw e;
+        }
+
+        return supportsTransactions = true;
+    }
+
     /**
      * Adds statement to openStatements
      * 
diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerConnectionPoolProxy.java b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerConnectionPoolProxy.java
index 5e4d933c..056c7f89 100644
--- a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerConnectionPoolProxy.java
+++ b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerConnectionPoolProxy.java
@@ -661,4 +661,25 @@ class SQLServerConnectionPoolProxy implements ISQLServerConnection, java.io.Seri
     @Deprecated(since = "12.1.0", forRemoval = true)
     @Override
     public void setMsiTokenCacheTtl(int timeToLive) {}
+
+    /**
+     * Returns the fully qualified class name of the implementing class for {@link SQLServerAccessTokenCallback}.
+     *
+     * @return accessTokenCallbackClass
+     */
+    @Override
+    public String getAccessTokenCallbackClass() {
+        return wrappedConnection.getAccessTokenCallbackClass();
+    }
+
+    /**
+     * Sets 'accessTokenCallbackClass' to the fully qualified class name
+     * of the implementing class for {@link SQLServerAccessTokenCallback}.
+     *
+     * @param accessTokenCallbackClass
+     */
+    @Override
+    public void setAccessTokenCallbackClass(String accessTokenCallbackClass) {
+        wrappedConnection.setAccessTokenCallbackClass(accessTokenCallbackClass);
+    }
 }
diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerDataSource.java b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerDataSource.java
index eeb9907a..47478035 100644
--- a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerDataSource.java
+++ b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerDataSource.java
@@ -1280,6 +1280,29 @@ public class SQLServerDataSource
                 SQLServerDriverObjectProperty.ACCESS_TOKEN_CALLBACK.getDefaultValue());
     }
 
+    /**
+     * Sets 'accessTokenCallbackClass' to the fully qualified class name
+     * of the implementing class for {@link SQLServerAccessTokenCallback}.
+     *
+     * @param accessTokenCallbackClass
+     */
+    @Override
+    public void setAccessTokenCallbackClass(String accessTokenCallbackClass) {
+        setStringProperty(connectionProps, SQLServerDriverStringProperty.ACCESS_TOKEN_CALLBACK_CLASS.toString(),
+                accessTokenCallbackClass);
+    }
+
+    /**
+     * Returns the fully qualified class name of the implementing class for {@link SQLServerAccessTokenCallback}.
+     *
+     * @return accessTokenCallbackClass
+     */
+    @Override
+    public String getAccessTokenCallbackClass() {
+        return getStringProperty(connectionProps, SQLServerDriverStringProperty.ACCESS_TOKEN_CALLBACK_CLASS.toString(),
+                null);
+    }
+
     /**
      * Sets a property string value.
      *
diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerDatabaseMetaData.java b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerDatabaseMetaData.java
index 17ad5a0b..79f2c269 100644
--- a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerDatabaseMetaData.java
+++ b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerDatabaseMetaData.java
@@ -1667,7 +1667,7 @@ public final class SQLServerDatabaseMetaData implements java.sql.DatabaseMetaDat
                     && !name.equals(SQLServerDriverStringProperty.TRUST_STORE_PASSWORD.toString())) {
                 String val = info[index].value;
                 // skip empty strings
-                if (0 != val.length()) {
+                if (null != val && 0 != val.length()) {
                     // special case these server name, instance name and port
                     // number as these go in the front
                     if (name.equals(SQLServerDriverStringProperty.SERVER_NAME.toString())) {
@@ -2171,7 +2171,7 @@ public final class SQLServerDatabaseMetaData implements java.sql.DatabaseMetaDat
     @Override
     public boolean supportsTransactions() throws SQLServerException {
         checkClosed();
-        return true;
+        return connection.supportsTransactions();
     }
 
     @Override
diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerDriver.java b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerDriver.java
index 95b0ebca..491eac24 100644
--- a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerDriver.java
+++ b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerDriver.java
@@ -605,7 +605,8 @@ enum SQLServerDriverStringProperty {
     MAX_RESULT_BUFFER("maxResultBuffer", "-1"),
     ENCRYPT("encrypt", EncryptOption.TRUE.toString()),
     SERVER_CERTIFICATE("serverCertificate", ""),
-    DATETIME_DATATYPE("datetimeParameterType", DatetimeType.DATETIME2.toString());
+    DATETIME_DATATYPE("datetimeParameterType", DatetimeType.DATETIME2.toString()),
+    ACCESS_TOKEN_CALLBACK_CLASS("accessTokenCallbackClass", "");
 
     private final String name;
     private final String defaultValue;
@@ -833,6 +834,12 @@ public final class SQLServerDriver implements java.sql.Driver {
                     SQLServerDriverStringProperty.TRUST_MANAGER_CLASS.getDefaultValue(), false, null),
             new SQLServerDriverPropertyInfo(SQLServerDriverStringProperty.TRUST_MANAGER_CONSTRUCTOR_ARG.toString(),
                     SQLServerDriverStringProperty.TRUST_MANAGER_CONSTRUCTOR_ARG.getDefaultValue(), false, null),
+            // Callback needs to be in list despite not being settable within connection string.
+            // The reason for this is for calls to mergeURLAndSuppliedProperties to work when setting the callback.
+            new SQLServerDriverPropertyInfo(SQLServerDriverObjectProperty.ACCESS_TOKEN_CALLBACK.toString(),
+                    SQLServerDriverObjectProperty.ACCESS_TOKEN_CALLBACK.getDefaultValue(), false, null),
+            new SQLServerDriverPropertyInfo(SQLServerDriverStringProperty.ACCESS_TOKEN_CALLBACK_CLASS.toString(),
+                    SQLServerDriverStringProperty.ACCESS_TOKEN_CALLBACK_CLASS.getDefaultValue(), false, null),
             new SQLServerDriverPropertyInfo(SQLServerDriverBooleanProperty.REPLICATION.toString(),
                     Boolean.toString(SQLServerDriverBooleanProperty.REPLICATION.getDefaultValue()), false, TRUE_FALSE),
             new SQLServerDriverPropertyInfo(SQLServerDriverBooleanProperty.SEND_TIME_AS_DATETIME.toString(),
@@ -931,6 +938,8 @@ public final class SQLServerDriver implements java.sql.Driver {
     private static final SQLServerDriverPropertyInfo[] DRIVER_PROPERTIES_PROPERTY_ONLY = {
             // default required available choices
             // property name value property (if appropriate)
+            new SQLServerDriverPropertyInfo(SQLServerDriverObjectProperty.ACCESS_TOKEN_CALLBACK.toString(),
+                    SQLServerDriverObjectProperty.ACCESS_TOKEN_CALLBACK.getDefaultValue(), false, null),
             new SQLServerDriverPropertyInfo(SQLServerDriverStringProperty.ACCESS_TOKEN.toString(),
                     SQLServerDriverStringProperty.ACCESS_TOKEN.getDefaultValue(), false, null),
             new SQLServerDriverPropertyInfo(SQLServerDriverObjectProperty.GSS_CREDENTIAL.toString(),
@@ -1074,6 +1083,9 @@ public final class SQLServerDriver implements java.sql.Driver {
                     fixedup.setProperty(newname, val);
                 } else if ("gsscredential".equalsIgnoreCase(newname) && (props.get(name) instanceof GSSCredential)) {
                     fixedup.put(newname, props.get(name));
+                } else if ("accessTokenCallback".equalsIgnoreCase(newname)
+                        && (props.get(name) instanceof SQLServerAccessTokenCallback)) {
+                    fixedup.put(newname, props.get(name));
                 } else {
                     MessageFormat form = new MessageFormat(SQLServerException.getErrString("R_invalidpropertyValue"));
                     Object[] msgArgs = {name};
diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerResource.java b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerResource.java
index 9f3901c4..c936a572 100644
--- a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerResource.java
+++ b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerResource.java
@@ -91,6 +91,7 @@ public final class SQLServerResource extends ListResourceBundle {
         {"R_invalidParameterNumber", "The parameter number {0} is not valid."},
         {"R_noMetadata", "There is no metadata."},
         {"R_resultsetClosed", "The result set is closed."},
+        {"R_transactionsNotSupported", "'@@TRANCOUNT' is not supported."},
         {"R_invalidColumnName", "The column name {0} is not valid."},
         {"R_resultsetNotUpdatable", "The result set is not updatable."},
         {"R_indexOutOfRange", "The index {0} is out of range."},
@@ -239,6 +240,8 @@ public final class SQLServerResource extends ListResourceBundle {
         {"R_delayLoadingLobsPropertyDescription", "Boolean flag which indicates whether the driver will load LOB datatypes into memory."},
         {"R_AADSecurePrincipalIdPropertyDescription", "The Application Id of a registered application which has been granted permission to the database connected."},
         {"R_AADSecurePrincipalSecretPropertyDescription", "A Secret defined for a registered application which has been granted permission to the database connected."},
+        {"R_accessTokenCallbackClassPropertyDescription", "The class to instantiate as the SQLServerAccessTokenCallback for acquiring tokens."},
+        {"R_accessTokenCallbackPropertyDescription", "A SQLServerAccessTokenCallback object which is used to call a callback method to return an access token."},
         {"R_noParserSupport", "An error occurred while instantiating the required parser. Error: \"{0}\""},
         {"R_writeOnlyXML", "Cannot read from this SQLXML instance. This instance is for writing data only."},
         {"R_dataHasBeenReadXML", "Cannot read from this SQLXML instance. The data has already been read."},
@@ -413,6 +416,7 @@ public final class SQLServerResource extends ListResourceBundle {
         {"R_AECertHashEmpty", "Internal error. Empty certificate thumbprint specified in certificate path {0}."},
         {"R_AECertNotFound", "Certificate with thumbprint {2} not found in certificate store {1} in certificate location {0}. Verify the certificate path in the column master key definition in the database is correct, and the certificate has been imported correctly into the certificate location/store."},
         {"R_AEMaloc", "Memory allocation failure."},
+        {"R_InvalidAccessTokenCallbackClass", "Invalid accessTokenCallbackClass: {0}"},
         {"R_AEKeypathLong", "Internal error. Specified certificate path has {0} bytes, which exceeds maximum length of {1} bytes."},
         {"R_AEECEKLenBad", "The specified encrypted column encryption key''s ciphertext length: {0} does not match the ciphertext length: {1} when using column master key (certificate) in \"{2}\". The encrypted column encryption key may be corrupt, or the specified certificate path may be incorrect."},
         {"R_AEECEKSigLenBad", "The specified encrypted column encryption key''s signature length {0} does not match the length {1} when using the column master key (certificate) in \"{2}\". The encrypted column encryption key may be corrupt, or the specified certificate path may be incorrect."},
diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/osgi/Activator.java b/src/main/java/com/microsoft/sqlserver/jdbc/osgi/Activator.java
index 241448ca..df29400e 100644
--- a/src/main/java/com/microsoft/sqlserver/jdbc/osgi/Activator.java
+++ b/src/main/java/com/microsoft/sqlserver/jdbc/osgi/Activator.java
@@ -38,6 +38,12 @@ public class Activator implements BundleActivator {
         properties.put(DataSourceFactory.OSGI_JDBC_DRIVER_NAME, "Microsoft JDBC Driver for SQL Server");
         properties.put(DataSourceFactory.OSGI_JDBC_DRIVER_VERSION,
                 driver.getMajorVersion() + "." + driver.getMinorVersion());
+        properties.put(DataSourceFactory.OSGI_JDBC_CAPABILITY, new String[] {
+                DataSourceFactory.OSGI_JDBC_CAPABILITY_DRIVER,
+                DataSourceFactory.OSGI_JDBC_CAPABILITY_DATASOURCE,
+                DataSourceFactory.OSGI_JDBC_CAPABILITY_CONNECTIONPOOLDATASOURCE,
+                DataSourceFactory.OSGI_JDBC_CAPABILITY_XADATASOURCE
+        });
         service = context.registerService(DataSourceFactory.class, new SQLServerDataSourceFactory(), properties);
         SQLServerDriver.register();
     }


Test patch:
diff --git a/src/test/java/com/microsoft/sqlserver/jdbc/SQLServerConnectionTest.java b/src/test/java/com/microsoft/sqlserver/jdbc/SQLServerConnectionTest.java
index 5b4b2a51..f1e2298c 100644
--- a/src/test/java/com/microsoft/sqlserver/jdbc/SQLServerConnectionTest.java
+++ b/src/test/java/com/microsoft/sqlserver/jdbc/SQLServerConnectionTest.java
@@ -4,7 +4,6 @@
  */
 package com.microsoft.sqlserver.jdbc;
 
-import static org.junit.Assert.assertNotEquals;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.junit.jupiter.api.Assertions.fail;
@@ -1108,68 +1107,4 @@ public class SQLServerConnectionTest extends AbstractTest {
             assertTrue(e.getMessage().matches(TestUtils.formatErrorMsg("R_errorServerName")));
         }
     }
-
-    @Tag(Constants.xSQLv11)
-    @Tag(Constants.xSQLv12)
-    @Tag(Constants.xSQLv14)
-    @Tag(Constants.xSQLv15)
-    @Tag(Constants.xAzureSQLDW)
-    @Tag(Constants.reqExternalSetup)
-    @Test
-    public void testDSPooledConnectionAccessTokenCallback() throws Exception {
-        SQLServerConnectionPoolDataSource ds = new SQLServerConnectionPoolDataSource();
-
-        // User/password is not required for access token callback
-        String cs = TestUtils.addOrOverrideProperty(connectionString, "user", "");
-        cs = TestUtils.addOrOverrideProperty(cs, "password", "");
-        AbstractTest.updateDataSource(cs, ds);
-        ds.setAccessTokenCallback(TestUtils.accessTokenCallback);
-
-        TestUtils.expireTokenToggle = false;
-        SQLServerPooledConnection pc = (SQLServerPooledConnection) ds.getPooledConnection();
-        String conn1ID;
-        String conn2ID;
-
-        // Callback should provide valid token on connection open for all new connections
-        // When the access token hasn't expired, the connection ID should be the same
-        try (Connection conn1 = pc.getConnection()) {}
-        conn1ID = TestUtils.getConnectionID(pc);
-        try (Connection conn2 = pc.getConnection()) {}
-        conn2ID = TestUtils.getConnectionID(pc);
-        assertEquals(conn1ID, conn2ID);
-    }
-
-    @Tag(Constants.xSQLv11)
-    @Tag(Constants.xSQLv12)
-    @Tag(Constants.xSQLv14)
-    @Tag(Constants.xSQLv15)
-    @Tag(Constants.xAzureSQLDW)
-    @Tag(Constants.reqExternalSetup)
-    @Test
-    public void testDSPooledConnectionAccessTokenCallbackExpiredToken() throws Exception {
-        SQLServerConnectionPoolDataSource ds = new SQLServerConnectionPoolDataSource();
-
-        // User/password is not required for access token callback
-        String cs = TestUtils.addOrOverrideProperty(connectionString, "user", "");
-        cs = TestUtils.addOrOverrideProperty(cs, "password", "");
-        AbstractTest.updateDataSource(cs, ds);
-        ds.setAccessTokenCallback(TestUtils.accessTokenCallback);
-
-        SQLServerPooledConnection pc = (SQLServerPooledConnection) ds.getPooledConnection();
-        String conn1ID;
-        String conn2ID;
-
-        // When token expires after first connection, it should create a new connection to get a new token.
-        // Connection ID should not be the same.
-        TestUtils.expireTokenToggle = true;
-        pc = (SQLServerPooledConnection) ds.getPooledConnection();
-        try (Connection conn1 = pc.getConnection()) {}
-        conn1ID = TestUtils.getConnectionID(pc);
-        // Sleep until token expires
-        Thread.sleep(720000);
-        try (Connection conn2 = pc.getConnection()) {}
-        conn2ID = TestUtils.getConnectionID(pc);
-
-        assertNotEquals(conn1ID, conn2ID);
-    }
 }
diff --git a/src/test/java/com/microsoft/sqlserver/jdbc/TestUtils.java b/src/test/java/com/microsoft/sqlserver/jdbc/TestUtils.java
index 88ad138f..2f20821a 100644
--- a/src/test/java/com/microsoft/sqlserver/jdbc/TestUtils.java
+++ b/src/test/java/com/microsoft/sqlserver/jdbc/TestUtils.java
@@ -96,6 +96,8 @@ public final class TestUtils {
 
     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;
 
@@ -132,9 +134,11 @@ public final class TestUtils {
                 String accessToken = authenticationResult.accessToken();
                 long expiresOn = authenticationResult.expiresOnDate().getTime();
 
+                ACCESS_TOKEN_CALLBACK = accessToken;
+
                 if (expireTokenToggle) {
                     Date now = new Date();
-                    long minutesToExpireWithin = 10 * 60 * 1000; // Expire within 10 minutes
+                    long minutesToExpireWithin = TEST_TOKEN_EXPIRY_SECONDS * 1000; // Expire within 2 minutes
                     return new SqlAuthenticationToken(accessToken, now.getTime() + minutesToExpireWithin);
                 } else {
                     return new SqlAuthenticationToken(accessToken, expiresOn);
@@ -161,6 +165,26 @@ public final class TestUtils {
         }
     }
 
+    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() {}
 
     /**
diff --git a/src/test/java/com/microsoft/sqlserver/jdbc/connection/RequestBoundaryMethodsTest.java b/src/test/java/com/microsoft/sqlserver/jdbc/connection/RequestBoundaryMethodsTest.java
index daba4459..d0ce2556 100644
--- a/src/test/java/com/microsoft/sqlserver/jdbc/connection/RequestBoundaryMethodsTest.java
+++ b/src/test/java/com/microsoft/sqlserver/jdbc/connection/RequestBoundaryMethodsTest.java
@@ -502,6 +502,8 @@ public class RequestBoundaryMethodsTest extends AbstractTest {
         verifiedMethodNames.add("setIPAddressPreference");
         verifiedMethodNames.add("getMsiTokenCacheTtl");
         verifiedMethodNames.add("setMsiTokenCacheTtl");
+        verifiedMethodNames.add("getAccessTokenCallbackClass");
+        verifiedMethodNames.add("setAccessTokenCallbackClass");
         return verifiedMethodNames;
     }
 }
diff --git a/src/test/java/com/microsoft/sqlserver/jdbc/fedauth/FedauthTest.java b/src/test/java/com/microsoft/sqlserver/jdbc/fedauth/FedauthTest.java
index 53c3d546..fc398312 100644
--- a/src/test/java/com/microsoft/sqlserver/jdbc/fedauth/FedauthTest.java
+++ b/src/test/java/com/microsoft/sqlserver/jdbc/fedauth/FedauthTest.java
@@ -40,9 +40,12 @@ public class FedauthTest extends FedauthCommon {
     static String charTable = TestUtils
             .escapeSingleQuotes(AbstractSQLGenerator.escapeIdentifier(RandomUtil.getIdentifier("JDBC_FedAuthTest")));
 
+    private static String accessTokenCallbackConnectionString;
+
     @BeforeAll
     public static void setupTests() throws Exception {
         connectionString = TestUtils.addOrOverrideProperty(connectionString, "trustServerCertificate", "true");
+        accessTokenCallbackConnectionString = "jdbc:sqlserver://" + azureServer + ";database=" + azureDatabase + ";";
         setConnection();
     }
 
@@ -352,6 +355,18 @@ public class FedauthTest extends FedauthCommon {
         validateException(url, "R_BothUserPasswordandDeprecated");
     }
 
+    @Tag(Constants.xSQLv11)
+    @Tag(Constants.xSQLv12)
+    @Tag(Constants.xSQLv14)
+    @Tag(Constants.xAzureSQLDW)
+    @Tag(Constants.reqExternalSetup)
+    @Test
+    public void testAccessTokenCallbackClassConnection() throws Exception {
+        String cs = TestUtils.addOrOverrideProperty(accessTokenCallbackConnectionString, "accessTokenCallbackClass",
+                PooledConnectionTest.AccessTokenCallbackClass.class.getName());
+        try (Connection conn1 = DriverManager.getConnection(cs)) {}
+    }
+
     private static void validateException(String url, String resourceKey) {
         try (Connection conn = DriverManager.getConnection(url)) {
             fail(TestResource.getResource("R_expectedFailPassed"));
diff --git a/src/test/java/com/microsoft/sqlserver/jdbc/fedauth/PooledConnectionTest.java b/src/test/java/com/microsoft/sqlserver/jdbc/fedauth/PooledConnectionTest.java
index 55415bbe..e109f0d6 100644
--- a/src/test/java/com/microsoft/sqlserver/jdbc/fedauth/PooledConnectionTest.java
+++ b/src/test/java/com/microsoft/sqlserver/jdbc/fedauth/PooledConnectionTest.java
@@ -4,6 +4,8 @@
  */
 package com.microsoft.sqlserver.jdbc.fedauth;
 
+import static org.junit.Assert.assertNotEquals;
+import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
 import static org.junit.jupiter.api.Assertions.fail;
 
@@ -13,9 +15,13 @@ import java.sql.DriverManager;
 import java.sql.SQLException;
 import java.sql.Statement;
 import java.util.ArrayList;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Random;
+import java.util.Set;
+import java.util.UUID;
 import java.util.concurrent.Callable;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
@@ -23,6 +29,16 @@ import java.util.concurrent.TimeUnit;
 
 import javax.sql.PooledConnection;
 
+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.jdbc.SQLServerAccessTokenCallback;
+import com.microsoft.sqlserver.jdbc.SQLServerException;
+import com.microsoft.sqlserver.jdbc.SqlAuthenticationToken;
+import com.microsoft.sqlserver.jdbc.TestResource;
+import com.microsoft.sqlserver.testframework.AbstractTest;
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Tag;
@@ -46,9 +62,40 @@ public class PooledConnectionTest extends FedauthCommon {
     static String charTable = TestUtils.escapeSingleQuotes(
             AbstractSQLGenerator.escapeIdentifier(RandomUtil.getIdentifier("JDBC_PooledConnection")));
 
+    private static String accessTokenCallbackConnectionString;
+
+    public static class AccessTokenCallbackClass implements 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();
+
+                return new SqlAuthenticationToken(accessToken, expiresOn);
+            } catch (Exception e) {
+                fail(TestResource.getResource("R_unexpectedException") + e.getMessage());
+            }
+            return null;
+        }
+    }
+
     @BeforeAll
     public static void setupTests() throws Exception {
         setConnection();
+        accessTokenCallbackConnectionString = "jdbc:sqlserver://" + azureServer + ";database=" + azureDatabase + ";";
     }
 
     @Test
@@ -263,6 +310,167 @@ public class PooledConnectionTest extends FedauthCommon {
         }
     }
 
+    @Tag(Constants.xSQLv11)
+    @Tag(Constants.xSQLv12)
+    @Tag(Constants.xSQLv14)
+    @Tag(Constants.xAzureSQLDW)
+    @Tag(Constants.reqExternalSetup)
+    @Test
+    public void testDSPooledConnectionAccessTokenCallback() throws Exception {
+        SQLServerConnectionPoolDataSource ds = new SQLServerConnectionPoolDataSource();
+
+        // User/password is not required for access token callback
+        AbstractTest.updateDataSource(accessTokenCallbackConnectionString, ds);
+        ds.setAccessTokenCallback(TestUtils.accessTokenCallback);
+
+        TestUtils.expireTokenToggle = false;
+        SQLServerPooledConnection pc = (SQLServerPooledConnection) ds.getPooledConnection();
+        String conn1ID;
+        String conn2ID;
+
+        // Callback should provide valid token on connection open for all new connections
+        // When the access token hasn't expired, the connection ID should be the same
+        try (Connection conn1 = pc.getConnection()) {}
+        conn1ID = TestUtils.getConnectionID(pc);
+        try (Connection conn2 = pc.getConnection()) {}
+        conn2ID = TestUtils.getConnectionID(pc);
+        assertEquals(conn1ID, conn2ID);
+    }
+
+    @Tag(Constants.xSQLv11)
+    @Tag(Constants.xSQLv12)
+    @Tag(Constants.xSQLv14)
+    @Tag(Constants.xAzureSQLDW)
+    @Tag(Constants.reqExternalSetup)
+    @Test
+    public void testDSPooledConnectionMergeAccessTokenCallbackProperty() throws Exception {
+        SQLServerConnectionPoolDataSource ds = new SQLServerConnectionPoolDataSource();
+
+        // Access token callback property from ds.setAccessTokenCallback should merge supplied properties from URL
+        ds.setURL(accessTokenCallbackConnectionString);
+        ds.setAccessTokenCallback(TestUtils.accessTokenCallback);
+
+        SQLServerPooledConnection pc = (SQLServerPooledConnection) ds.getPooledConnection();
+        String conn1ID;
+        String conn2ID;
+
+        try (Connection conn1 = pc.getConnection()) {}
+        conn1ID = TestUtils.getConnectionID(pc);
+        try (Connection conn2 = pc.getConnection()) {}
+        conn2ID = TestUtils.getConnectionID(pc);
+        assertEquals(conn1ID, conn2ID);
+    }
+
+    @Tag(Constants.xSQLv11)
+    @Tag(Constants.xSQLv12)
+    @Tag(Constants.xSQLv14)
+    @Tag(Constants.xAzureSQLDW)
+    @Tag(Constants.reqExternalSetup)
+    @Test
+    public void testDSPooledConnectionAccessTokenCallbackExpiredToken() throws Exception {
+        SQLServerConnectionPoolDataSource ds = new SQLServerConnectionPoolDataSource();
+
+        // User/password is not required for access token callback
+        AbstractTest.updateDataSource(accessTokenCallbackConnectionString, ds);
+        ds.setAccessTokenCallback(TestUtils.accessTokenCallback);
+
+        SQLServerPooledConnection pc = (SQLServerPooledConnection) ds.getPooledConnection();
+        String conn1ID;
+        String conn2ID;
+
+        // When token expires after first connection, it should create a new connection to get a new token.
+        // Connection ID should not be the same.
+        TestUtils.expireTokenToggle = true;
+        pc = (SQLServerPooledConnection) ds.getPooledConnection();
+        try (Connection conn1 = pc.getConnection()) {
+            TestUtils.setAccessTokenExpiry(conn1);
+        }
+        conn1ID = TestUtils.getConnectionID(pc);
+        // Sleep until token expires
+        Thread.sleep(130 * 1000);
+        try (Connection conn2 = pc.getConnection()) {}
+        conn2ID = TestUtils.getConnectionID(pc);
+
+        assertNotEquals(conn1ID, conn2ID);
+    }
+
+    @Tag(Constants.xSQLv11)
+    @Tag(Constants.xSQLv12)
+    @Tag(Constants.xSQLv14)
+    @Tag(Constants.xAzureSQLDW)
+    @Tag(Constants.reqExternalSetup)
+    @Test
+    public void testDSPooledConnectionAccessTokenCallbackClass() throws Exception {
+        SQLServerConnectionPoolDataSource ds = new SQLServerConnectionPoolDataSource();
+
+        // User/password is not required for access token callback
+        AbstractTest.updateDataSource(accessTokenCallbackConnectionString, ds);
+        ds.setAccessTokenCallbackClass(AccessTokenCallbackClass.class.getName());
+
+        SQLServerPooledConnection pc = (SQLServerPooledConnection) ds.getPooledConnection();
+        String conn1ID;
+        String conn2ID;
+
+        // Callback should provide valid token on connection open for all new connections
+        // When the access token hasn't expired, the connection ID should be the same
+        try (Connection conn1 = pc.getConnection()) {}
+        conn1ID = TestUtils.getConnectionID(pc);
+        try (Connection conn2 = pc.getConnection()) {}
+        conn2ID = TestUtils.getConnectionID(pc);
+        assertEquals(conn1ID, conn2ID);
+    }
+
+    @Tag(Constants.xSQLv11)
+    @Tag(Constants.xSQLv12)
+    @Tag(Constants.xSQLv14)
+    @Tag(Constants.xAzureSQLDW)
+    @Tag(Constants.reqExternalSetup)
+    @Test
+    public void testDSPooledConnectionAccessTokenCallbackClassExceptions() throws Exception {
+        SQLServerConnectionPoolDataSource ds = new SQLServerConnectionPoolDataSource();
+
+        // User/password is not required for access token callback
+        AbstractTest.updateDataSource(accessTokenCallbackConnectionString, ds);
+        ds.setAccessTokenCallbackClass(AccessTokenCallbackClass.class.getName());
+        ds.setUser("user");
+        SQLServerPooledConnection pc;
+
+        // Should fail with user set
+        try {
+            pc = (SQLServerPooledConnection) ds.getPooledConnection();
+            fail(TestResource.getResource("R_expectedFailPassed"));
+        } catch (SQLServerException e) {
+            assertTrue(e.getMessage().matches(TestUtils.formatErrorMsg("R_AccessTokenCallbackWithUserPassword")));
+        }
+
+        ds.setUser("");
+        ds.setPassword(UUID.randomUUID().toString());
+
+        // Should fail with password set
+        try {
+            pc = (SQLServerPooledConnection) ds.getPooledConnection();
+            fail(TestResource.getResource("R_expectedFailPassed"));
+        } catch (SQLServerException e) {
+            assertTrue(e.getMessage().matches(TestUtils.formatErrorMsg("R_AccessTokenCallbackWithUserPassword")));
+        }
+
+        // Should fail with invalid accessTokenCallbackClass value
+        ds.setAccessTokenCallbackClass("Invalid");
+        ds.setUser("");
+        ds.setPassword("");
+        try {
+            pc = (SQLServerPooledConnection) ds.getPooledConnection();
+            fail(TestResource.getResource("R_expectedFailPassed"));
+        } catch (SQLServerException e) {
+            assertTrue(e.getMessage().matches(TestUtils.formatErrorMsg("R_InvalidAccessTokenCallbackClass")));
+        }
+
+        // Should pass with no user or password set
+        ds.setAccessTokenCallbackClass(AccessTokenCallbackClass.class.getName());
+        pc = (SQLServerPooledConnection) ds.getPooledConnection();
+        try (Connection conn1 = pc.getConnection()) {}
+    }
+
     @AfterAll
     public static void terminate() throws SQLException {
         try (Connection conn = DriverManager.getConnection(adPasswordConnectionStr);
diff --git a/src/test/java/com/microsoft/sqlserver/jdbc/osgi/DataFactoryTest.java b/src/test/java/com/microsoft/sqlserver/jdbc/osgi/DataFactoryTest.java
index bfa8c3fe..62731692 100644
--- a/src/test/java/com/microsoft/sqlserver/jdbc/osgi/DataFactoryTest.java
+++ b/src/test/java/com/microsoft/sqlserver/jdbc/osgi/DataFactoryTest.java
@@ -10,8 +10,11 @@ import java.sql.Connection;
 import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.Statement;
+import java.util.Arrays;
 import java.util.Dictionary;
+import java.util.HashSet;
 import java.util.Properties;
+import java.util.Set;
 
 import javax.sql.ConnectionPoolDataSource;
 import javax.sql.DataSource;
@@ -19,7 +22,6 @@ import javax.sql.PooledConnection;
 import javax.sql.XAConnection;
 import javax.sql.XADataSource;
 
-import com.microsoft.sqlserver.jdbc.TestUtils;
 import org.eclipse.gemini.blueprint.mock.MockBundleContext;
 import org.eclipse.gemini.blueprint.mock.MockServiceRegistration;
 import org.junit.jupiter.api.BeforeAll;
@@ -87,6 +89,12 @@ public class DataFactoryTest extends AbstractTest {
         boolean correctClass = false;
         boolean correctName = false;
         boolean correctVersion = false;
+        String[] required_capabilities=new String[] {
+                DataSourceFactory.OSGI_JDBC_CAPABILITY_DRIVER,
+                DataSourceFactory.OSGI_JDBC_CAPABILITY_DATASOURCE,
+                DataSourceFactory.OSGI_JDBC_CAPABILITY_CONNECTIONPOOLDATASOURCE,
+                DataSourceFactory.OSGI_JDBC_CAPABILITY_XADATASOURCE};
+        boolean driverCap = false;
         SQLServerDriver driver = new SQLServerDriver();
 
         for (String key : propertyKeys) {
@@ -110,9 +118,18 @@ public class DataFactoryTest extends AbstractTest {
                         "Driver version mismatch. Expected: " + bundleDriverVer + ", Actual: " + actualDriverVer + ".",
                         bundleDriverVer.equals(actualDriverVer));
                 correctVersion = true;
+            }else if (key.equals(DataSourceFactory.OSGI_JDBC_CAPABILITY)) {
+                Object property = sr.getProperty(key);
+                if (property instanceof String[]) {
+                    Set<String> set = new HashSet<>(Arrays.asList( (String[]) property));
+                    for (String cap : required_capabilities) {
+                        assertTrue("Capability "+cap+" is missing!", set.contains(cap));
+                    }
+                }
+                driverCap=true;
             }
         }
-        assertTrue("Not all properties were checked.", correctClass && correctName && correctVersion);
+        assertTrue("Not all properties were checked.", correctClass && correctName && correctVersion && driverCap);
         DataSourceFactory dsf = bc.getService(sr);
         verifyFactoryNormalConnection(dsf);
         verifyFactoryPooledConnection(dsf);