Add | Add Workload Identity Support (#2159)

dotnet · Dec 7, 2023 · 17590d3 · 17590d3
1 parent 4319661
commit 17590d3
Show file tree

Hide file tree

Showing 20 changed files with 185 additions and 25 deletions.
diff --git a/doc/snippets/Microsoft.Data.SqlClient/SqlAuthenticationMethod.xml b/doc/snippets/Microsoft.Data.SqlClient/SqlAuthenticationMethod.xml
@@ -45,5 +45,9 @@
           <summary>The authentication method uses Active Directory Default. Use this mode to connect to a SQL Database using multiple non-interactive authentication methods tried sequentially to acquire an access token. This method does not fallback to the "Active Directory Interactive" authentication method.</summary>
           <value>9</value>
       </ActiveDirectoryDefault>
+      <ActiveDirectoryWorkloadIdentity>
+        <summary>The authentication method uses Active Directory Workload Identity. Use a federated User Assigned Managed Identity to connect to SQL Database from Azure client environments that have enabled support for Workload Identity. The 'User Id' or 'UID' is required to be set to the "client ID" of the user identity.</summary>
+        <value>10</value>
+      </ActiveDirectoryWorkloadIdentity>
     </members>
 </docs>
diff --git a/src/Microsoft.Data.SqlClient.sln b/src/Microsoft.Data.SqlClient.sln
@@ -11,9 +11,6 @@ EndProject
 Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "TDS.Servers", "Microsoft.Data.SqlClient\tests\tools\TDS\TDS.Servers\TDS.Servers.csproj", "{978063D3-FBB5-4E10-8C45-67C90BE1B931}"
 EndProject
 Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "TDS", "Microsoft.Data.SqlClient\tests\tools\TDS\TDS\TDS.csproj", "{8DC9D1A0-351B-47BC-A90F-B9DA542550E9}"
-	ProjectSection(ProjectDependencies) = postProject
-		{FDA6971D-9F57-4DA4-B10A-261C91684CFC} = {FDA6971D-9F57-4DA4-B10A-261C91684CFC}
-	EndProjectSection
 EndProject
 Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "Microsoft.Data.SqlClient.Tests", "Microsoft.Data.SqlClient\tests\FunctionalTests\Microsoft.Data.SqlClient.Tests.csproj", "{D2D1E2D1-B6E0-489F-A36D-1F3047AB87B9}"
 	ProjectSection(ProjectDependencies) = postProject
@@ -28,7 +25,6 @@ Project("{2150E333-8FDC-42A3-9474-1A3956D46DE8}") = "netfx", "netfx", "{3FDD425C
 EndProject
 Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "Address", "Microsoft.Data.SqlClient\tests\ManualTests\SQL\UdtTest\UDTs\Address\Address.csproj", "{D1392B54-998A-4F27-BC17-4CE149117BCC}"
 	ProjectSection(ProjectDependencies) = postProject
-		{FDA6971D-9F57-4DA4-B10A-261C91684CFC} = {FDA6971D-9F57-4DA4-B10A-261C91684CFC}
 		{37431336-5307-4184-9356-C4B7E47DC714} = {37431336-5307-4184-9356-C4B7E47DC714}
 		{407890AC-9876-4FEF-A6F1-F36A876BAADE} = {407890AC-9876-4FEF-A6F1-F36A876BAADE}
 	EndProjectSection
@@ -40,13 +36,11 @@ Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "Microsoft.Data.SqlClient.Ma
 EndProject
 Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "Circle", "Microsoft.Data.SqlClient\tests\ManualTests\SQL\UdtTest\UDTs\Circle\Circle.csproj", "{6C88F00F-9597-43AD-9E5F-9B344DA3B16F}"
 	ProjectSection(ProjectDependencies) = postProject
-		{FDA6971D-9F57-4DA4-B10A-261C91684CFC} = {FDA6971D-9F57-4DA4-B10A-261C91684CFC}
 		{37431336-5307-4184-9356-C4B7E47DC714} = {37431336-5307-4184-9356-C4B7E47DC714}
 	EndProjectSection
 EndProject
 Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "Shapes", "Microsoft.Data.SqlClient\tests\ManualTests\SQL\UdtTest\UDTs\Shapes\Shapes.csproj", "{B73A7063-37C3-415D-AD53-BB3DA20ABD6E}"
 	ProjectSection(ProjectDependencies) = postProject
-		{FDA6971D-9F57-4DA4-B10A-261C91684CFC} = {FDA6971D-9F57-4DA4-B10A-261C91684CFC}
 		{37431336-5307-4184-9356-C4B7E47DC714} = {37431336-5307-4184-9356-C4B7E47DC714}
 	EndProjectSection
 EndProject
@@ -384,18 +378,6 @@ Global
 		{9073ABEF-92E0-4702-BB23-2C99CEF9BDD7}.Release|x64.Build.0 = Release|Any CPU
 		{9073ABEF-92E0-4702-BB23-2C99CEF9BDD7}.Release|x86.ActiveCfg = Release|Any CPU
 		{9073ABEF-92E0-4702-BB23-2C99CEF9BDD7}.Release|x86.Build.0 = Release|Any CPU
-		{FDA6971D-9F57-4DA4-B10A-261C91684CFC}.Debug|Any CPU.ActiveCfg = Debug|Any CPU
-		{FDA6971D-9F57-4DA4-B10A-261C91684CFC}.Debug|Any CPU.Build.0 = Debug|Any CPU
-		{FDA6971D-9F57-4DA4-B10A-261C91684CFC}.Debug|x64.ActiveCfg = Debug|Any CPU
-		{FDA6971D-9F57-4DA4-B10A-261C91684CFC}.Debug|x64.Build.0 = Debug|Any CPU
-		{FDA6971D-9F57-4DA4-B10A-261C91684CFC}.Debug|x86.ActiveCfg = Debug|Any CPU
-		{FDA6971D-9F57-4DA4-B10A-261C91684CFC}.Debug|x86.Build.0 = Debug|Any CPU
-		{FDA6971D-9F57-4DA4-B10A-261C91684CFC}.Release|Any CPU.ActiveCfg = Release|Any CPU
-		{FDA6971D-9F57-4DA4-B10A-261C91684CFC}.Release|Any CPU.Build.0 = Release|Any CPU
-		{FDA6971D-9F57-4DA4-B10A-261C91684CFC}.Release|x64.ActiveCfg = Release|Any CPU
-		{FDA6971D-9F57-4DA4-B10A-261C91684CFC}.Release|x64.Build.0 = Release|Any CPU
-		{FDA6971D-9F57-4DA4-B10A-261C91684CFC}.Release|x86.ActiveCfg = Release|Any CPU
-		{FDA6971D-9F57-4DA4-B10A-261C91684CFC}.Release|x86.Build.0 = Release|Any CPU
 		{F5DF2FDC-C860-4CB3-8B24-7C903C6FC076}.Debug|Any CPU.ActiveCfg = Debug|Any CPU
 		{F5DF2FDC-C860-4CB3-8B24-7C903C6FC076}.Debug|x64.ActiveCfg = Debug|Any CPU
 		{F5DF2FDC-C860-4CB3-8B24-7C903C6FC076}.Debug|x64.Build.0 = Debug|Any CPU
@@ -511,7 +493,6 @@ Global
 		{771F3F1E-7A68-4A9D-ADA8-A24F1D5BE71D} = {3FDD425C-FE01-4B56-863E-1FCDD0677CF5}
 		{412BCCC8-19F6-489A-B594-E9A506816155} = {771F3F1E-7A68-4A9D-ADA8-A24F1D5BE71D}
 		{9073ABEF-92E0-4702-BB23-2C99CEF9BDD7} = {C9726AED-D6A3-4AAC-BA04-92DD1F079594}
-		{FDA6971D-9F57-4DA4-B10A-261C91684CFC} = {0CC4817A-12F3-4357-912C-09315FAAD008}
 		{71F356DC-DFA3-4163-8BFE-D268722CE189} = {ED952CF7-84DF-437A-B066-F516E9BE1C2C}
 		{908C7DD3-C999-40A6-9433-9F5ACA7C36F5} = {71F356DC-DFA3-4163-8BFE-D268722CE189}
 		{0CE216CE-8072-4985-B248-61F0D0BE9C2E} = {71F356DC-DFA3-4163-8BFE-D268722CE189}

diff --git a/...lient/netcore/src/Microsoft/Data/SqlClient/SqlAuthenticationProviderManager.NetCoreApp.cs b/...lient/netcore/src/Microsoft/Data/SqlClient/SqlAuthenticationProviderManager.NetCoreApp.cs
@@ -153,6 +153,8 @@ private static SqlAuthenticationMethod AuthenticationEnumFromString(string authe
                     return SqlAuthenticationMethod.ActiveDirectoryMSI;
                 case ActiveDirectoryDefault:
                     return SqlAuthenticationMethod.ActiveDirectoryDefault;
+                case ActiveDirectoryWorkloadIdentity:
+                    return SqlAuthenticationMethod.ActiveDirectoryWorkloadIdentity;
                 default:
                     throw SQL.UnsupportedAuthentication(authentication);
             }

diff --git a/...t.Data.SqlClient/netcore/src/Microsoft/Data/SqlClient/SqlAuthenticationProviderManager.cs b/...t.Data.SqlClient/netcore/src/Microsoft/Data/SqlClient/SqlAuthenticationProviderManager.cs
@@ -20,6 +20,7 @@ internal partial class SqlAuthenticationProviderManager
         private const string ActiveDirectoryManagedIdentity = "active directory managed identity";
         private const string ActiveDirectoryMSI = "active directory msi";
         private const string ActiveDirectoryDefault = "active directory default";
+        private const string ActiveDirectoryWorkloadIdentity = "active directory workload identity";
 
         private readonly IReadOnlyCollection<SqlAuthenticationMethod> _authenticationsWithAppSpecifiedProvider;
         private readonly ConcurrentDictionary<SqlAuthenticationMethod, SqlAuthenticationProvider> _providers;
@@ -45,6 +46,7 @@ private static void SetDefaultAuthProviders(SqlAuthenticationProviderManager ins
                 instance.SetProvider(SqlAuthenticationMethod.ActiveDirectoryManagedIdentity, activeDirectoryAuthProvider);
                 instance.SetProvider(SqlAuthenticationMethod.ActiveDirectoryMSI, activeDirectoryAuthProvider);
                 instance.SetProvider(SqlAuthenticationMethod.ActiveDirectoryDefault, activeDirectoryAuthProvider);
+                instance.SetProvider(SqlAuthenticationMethod.ActiveDirectoryWorkloadIdentity, activeDirectoryAuthProvider);
             }
         }
         /// <summary>

diff --git a/src/Microsoft.Data.SqlClient/netcore/src/Microsoft/Data/SqlClient/SqlConnection.cs b/src/Microsoft.Data.SqlClient/netcore/src/Microsoft/Data/SqlClient/SqlConnection.cs
@@ -202,6 +202,10 @@ public SqlConnection(string connectionString, SqlCredential credential) : this()
                 {
                     throw SQL.SettingCredentialWithNonInteractiveArgument(DbConnectionStringBuilderUtil.ActiveDirectoryDefaultString);
                 }
+                else if (UsesActiveDirectoryWorkloadIdentity(connectionOptions))
+                {
+                    throw SQL.SettingCredentialWithNonInteractiveArgument(DbConnectionStringBuilderUtil.ActiveDirectoryWorkloadIdentityString);
+                }
 
                 Credential = credential;
             }
@@ -530,6 +534,11 @@ private bool UsesActiveDirectoryDefault(SqlConnectionString opt)
             return opt != null && opt.Authentication == SqlAuthenticationMethod.ActiveDirectoryDefault;
         }
 
+        private bool UsesActiveDirectoryWorkloadIdentity(SqlConnectionString opt)
+        {
+            return opt != null && opt.Authentication == SqlAuthenticationMethod.ActiveDirectoryWorkloadIdentity;
+        }
+
         private bool UsesAuthentication(SqlConnectionString opt)
         {
             return opt != null && opt.Authentication != SqlAuthenticationMethod.NotSpecified;
@@ -619,6 +628,10 @@ public override string ConnectionString
                         {
                             throw SQL.SettingNonInteractiveWithCredential(DbConnectionStringBuilderUtil.ActiveDirectoryDefaultString);
                         }
+                        else if (UsesActiveDirectoryWorkloadIdentity(connectionOptions))
+                        {
+                            throw SQL.SettingNonInteractiveWithCredential(DbConnectionStringBuilderUtil.ActiveDirectoryWorkloadIdentityString);
+                        }
 
                         CheckAndThrowOnInvalidCombinationOfConnectionStringAndSqlCredential(connectionOptions);
                     }
@@ -999,6 +1012,10 @@ public SqlCredential Credential
                     {
                         throw SQL.SettingCredentialWithNonInteractiveInvalid(DbConnectionStringBuilderUtil.ActiveDirectoryDefaultString);
                     }
+                    else if (UsesActiveDirectoryWorkloadIdentity(connectionOptions))
+                    {
+                        throw SQL.SettingCredentialWithNonInteractiveInvalid(DbConnectionStringBuilderUtil.ActiveDirectoryWorkloadIdentityString);
+                    }
 
                     CheckAndThrowOnInvalidCombinationOfConnectionStringAndSqlCredential(connectionOptions);
 

diff --git a/...Microsoft.Data.SqlClient/netcore/src/Microsoft/Data/SqlClient/SqlInternalConnectionTds.cs b/...Microsoft.Data.SqlClient/netcore/src/Microsoft/Data/SqlClient/SqlInternalConnectionTds.cs
@@ -1330,6 +1330,7 @@ private void Login(ServerInfo server, TimeoutTimer timeout, string newPassword,
                 || ConnectionOptions.Authentication == SqlAuthenticationMethod.ActiveDirectoryManagedIdentity
                 || ConnectionOptions.Authentication == SqlAuthenticationMethod.ActiveDirectoryMSI
                 || ConnectionOptions.Authentication == SqlAuthenticationMethod.ActiveDirectoryDefault
+                || ConnectionOptions.Authentication == SqlAuthenticationMethod.ActiveDirectoryWorkloadIdentity
                 // Since AD Integrated may be acting like Windows integrated, additionally check _fedAuthRequired
                 || (ConnectionOptions.Authentication == SqlAuthenticationMethod.ActiveDirectoryIntegrated && _fedAuthRequired)
                 || _accessTokenCallback != null)
@@ -2159,6 +2160,7 @@ internal void OnFedAuthInfo(SqlFedAuthInfo fedAuthInfo)
                          || ConnectionOptions.Authentication == SqlAuthenticationMethod.ActiveDirectoryManagedIdentity
                          || ConnectionOptions.Authentication == SqlAuthenticationMethod.ActiveDirectoryMSI
                          || ConnectionOptions.Authentication == SqlAuthenticationMethod.ActiveDirectoryDefault
+                         || ConnectionOptions.Authentication == SqlAuthenticationMethod.ActiveDirectoryWorkloadIdentity
                          || (ConnectionOptions.Authentication == SqlAuthenticationMethod.ActiveDirectoryIntegrated && _fedAuthRequired),
                          "Credentials aren't provided for calling MSAL");
             Debug.Assert(fedAuthInfo != null, "info should not be null.");
@@ -2406,6 +2408,7 @@ internal SqlFedAuthToken GetFedAuthToken(SqlFedAuthInfo fedAuthInfo)
                         case SqlAuthenticationMethod.ActiveDirectoryManagedIdentity:
                         case SqlAuthenticationMethod.ActiveDirectoryMSI:
                         case SqlAuthenticationMethod.ActiveDirectoryDefault:
+                        case SqlAuthenticationMethod.ActiveDirectoryWorkloadIdentity:
                             if (_activeDirectoryAuthTimeoutRetryHelper.State == ActiveDirectoryAuthenticationTimeoutRetryState.Retrying)
                             {
                                 _fedAuthToken = _activeDirectoryAuthTimeoutRetryHelper.CachedToken;

diff --git a/src/Microsoft.Data.SqlClient/netcore/src/Microsoft/Data/SqlClient/TdsParser.cs b/src/Microsoft.Data.SqlClient/netcore/src/Microsoft/Data/SqlClient/TdsParser.cs
@@ -7973,6 +7973,9 @@ internal int WriteFedAuthFeatureRequest(FederatedAuthenticationFeatureExtensionD
                             case SqlAuthenticationMethod.ActiveDirectoryDefault:
                                 workflow = TdsEnums.MSALWORKFLOW_ACTIVEDIRECTORYDEFAULT;
                                 break;
+                            case SqlAuthenticationMethod.ActiveDirectoryWorkloadIdentity:
+                                workflow = TdsEnums.MSALWORKFLOW_ACTIVEDIRECTORYWORKLOADIDENTITY;
+                                break;
                             default:
                                 if (_connHandler._accessTokenCallback != null)
                                 {

diff --git a/...oft.Data.SqlClient/netfx/src/Microsoft/Data/SqlClient/SqlAuthenticationProviderManager.cs b/...oft.Data.SqlClient/netfx/src/Microsoft/Data/SqlClient/SqlAuthenticationProviderManager.cs
@@ -22,6 +22,7 @@ internal class SqlAuthenticationProviderManager
         private const string ActiveDirectoryManagedIdentity = "active directory managed identity";
         private const string ActiveDirectoryMSI = "active directory msi";
         private const string ActiveDirectoryDefault = "active directory default";
+        private const string ActiveDirectoryWorkloadIdentity = "active directory workload identity";
 
         static SqlAuthenticationProviderManager()
         {
@@ -52,6 +53,7 @@ static SqlAuthenticationProviderManager()
             Instance.SetProvider(SqlAuthenticationMethod.ActiveDirectoryManagedIdentity, activeDirectoryAuthProvider);
             Instance.SetProvider(SqlAuthenticationMethod.ActiveDirectoryMSI, activeDirectoryAuthProvider);
             Instance.SetProvider(SqlAuthenticationMethod.ActiveDirectoryDefault, activeDirectoryAuthProvider);
+            Instance.SetProvider(SqlAuthenticationMethod.ActiveDirectoryWorkloadIdentity, activeDirectoryAuthProvider);
         }
         public static readonly SqlAuthenticationProviderManager Instance;
 
@@ -231,6 +233,8 @@ private static SqlAuthenticationMethod AuthenticationEnumFromString(string authe
                     return SqlAuthenticationMethod.ActiveDirectoryMSI;
                 case ActiveDirectoryDefault:
                     return SqlAuthenticationMethod.ActiveDirectoryDefault;
+                case ActiveDirectoryWorkloadIdentity:
+                    return SqlAuthenticationMethod.ActiveDirectoryWorkloadIdentity;
                 default:
                     throw SQL.UnsupportedAuthentication(authentication);
             }

diff --git a/src/Microsoft.Data.SqlClient/netfx/src/Microsoft/Data/SqlClient/SqlConnection.cs b/src/Microsoft.Data.SqlClient/netfx/src/Microsoft/Data/SqlClient/SqlConnection.cs
@@ -405,6 +405,11 @@ public SqlConnection(string connectionString, SqlCredential credential) : this()
                     throw SQL.SettingCredentialWithNonInteractiveArgument(DbConnectionStringBuilderUtil.ActiveDirectoryDefaultString);
                 }
 
+                if (UsesActiveDirectoryWorkloadIdentity(connectionOptions))
+                {
+                    throw SQL.SettingCredentialWithNonInteractiveArgument(DbConnectionStringBuilderUtil.ActiveDirectoryWorkloadIdentityString);
+                }
+
                 Credential = credential;
             }
             // else
@@ -627,6 +632,11 @@ private bool UsesActiveDirectoryDefault(SqlConnectionString opt)
             return opt != null && opt.Authentication == SqlAuthenticationMethod.ActiveDirectoryDefault;
         }
 
+        private bool UsesActiveDirectoryWorkloadIdentity(SqlConnectionString opt)
+        {
+            return opt != null && opt.Authentication == SqlAuthenticationMethod.ActiveDirectoryWorkloadIdentity;
+        }
+
         private bool UsesAuthentication(SqlConnectionString opt)
         {
             return opt != null && opt.Authentication != SqlAuthenticationMethod.NotSpecified;
@@ -834,6 +844,10 @@ override public string ConnectionString
                         {
                             throw SQL.SettingNonInteractiveWithCredential(DbConnectionStringBuilderUtil.ActiveDirectoryDefaultString);
                         }
+                        else if (UsesActiveDirectoryWorkloadIdentity(connectionOptions))
+                        {
+                            throw SQL.SettingNonInteractiveWithCredential(DbConnectionStringBuilderUtil.ActiveDirectoryWorkloadIdentityString);
+                        }
 
                         CheckAndThrowOnInvalidCombinationOfConnectionStringAndSqlCredential(connectionOptions);
                     }
@@ -1183,6 +1197,10 @@ public SqlCredential Credential
                     {
                         throw SQL.SettingCredentialWithNonInteractiveInvalid(DbConnectionStringBuilderUtil.ActiveDirectoryDefaultString);
                     }
+                    else if (UsesActiveDirectoryWorkloadIdentity(connectionOptions))
+                    {
+                        throw SQL.SettingCredentialWithNonInteractiveInvalid(DbConnectionStringBuilderUtil.ActiveDirectoryWorkloadIdentityString);
+                    }
 
                     CheckAndThrowOnInvalidCombinationOfConnectionStringAndSqlCredential(connectionOptions);
 

diff --git a/src/Microsoft.Data.SqlClient/netfx/src/Microsoft/Data/SqlClient/SqlInternalConnectionTds.cs b/src/Microsoft.Data.SqlClient/netfx/src/Microsoft/Data/SqlClient/SqlInternalConnectionTds.cs
@@ -1595,6 +1595,7 @@ private void Login(ServerInfo server, TimeoutTimer timeout, string newPassword,
                 || ConnectionOptions.Authentication == SqlAuthenticationMethod.ActiveDirectoryManagedIdentity
                 || ConnectionOptions.Authentication == SqlAuthenticationMethod.ActiveDirectoryMSI
                 || ConnectionOptions.Authentication == SqlAuthenticationMethod.ActiveDirectoryDefault
+                || ConnectionOptions.Authentication == SqlAuthenticationMethod.ActiveDirectoryWorkloadIdentity
                 // Since AD Integrated may be acting like Windows integrated, additionally check _fedAuthRequired
                 || (ConnectionOptions.Authentication == SqlAuthenticationMethod.ActiveDirectoryIntegrated && _fedAuthRequired)
                 || _accessTokenCallback != null)
@@ -1991,7 +1992,8 @@ private bool ShouldDisableTnir(SqlConnectionString connectionOptions)
                                        connectionOptions.Authentication == SqlAuthenticationMethod.ActiveDirectoryDeviceCodeFlow ||
                                        connectionOptions.Authentication == SqlAuthenticationMethod.ActiveDirectoryManagedIdentity ||
                                        connectionOptions.Authentication == SqlAuthenticationMethod.ActiveDirectoryMSI ||
-                                       connectionOptions.Authentication == SqlAuthenticationMethod.ActiveDirectoryDefault;
+                                       connectionOptions.Authentication == SqlAuthenticationMethod.ActiveDirectoryDefault ||
+                                       connectionOptions.Authentication == SqlAuthenticationMethod.ActiveDirectoryWorkloadIdentity;
 
             // Check if the user had explicitly specified the TNIR option in the connection string or the connection string builder.
             // If the user has specified the option in the connection string explicitly, then we shouldn't disable TNIR.
@@ -2585,6 +2587,7 @@ internal void OnFedAuthInfo(SqlFedAuthInfo fedAuthInfo)
                          || ConnectionOptions.Authentication == SqlAuthenticationMethod.ActiveDirectoryManagedIdentity
                          || ConnectionOptions.Authentication == SqlAuthenticationMethod.ActiveDirectoryMSI
                          || ConnectionOptions.Authentication == SqlAuthenticationMethod.ActiveDirectoryDefault
+                         || ConnectionOptions.Authentication == SqlAuthenticationMethod.ActiveDirectoryWorkloadIdentity
                          || ConnectionOptions.Authentication == SqlAuthenticationMethod.ActiveDirectoryDeviceCodeFlow
                          || (ConnectionOptions.Authentication == SqlAuthenticationMethod.ActiveDirectoryIntegrated && _fedAuthRequired),
                          "Credentials aren't provided for calling MSAL");
@@ -2820,6 +2823,7 @@ internal SqlFedAuthToken GetFedAuthToken(SqlFedAuthInfo fedAuthInfo)
                         case SqlAuthenticationMethod.ActiveDirectoryManagedIdentity:
                         case SqlAuthenticationMethod.ActiveDirectoryMSI:
                         case SqlAuthenticationMethod.ActiveDirectoryDefault:
+                        case SqlAuthenticationMethod.ActiveDirectoryWorkloadIdentity:
                             if (_activeDirectoryAuthTimeoutRetryHelper.State == ActiveDirectoryAuthenticationTimeoutRetryState.Retrying)
                             {
                                 _fedAuthToken = _activeDirectoryAuthTimeoutRetryHelper.CachedToken;