From d84508b8580eabe1b0bc6f66015cf096a64eb988 Mon Sep 17 00:00:00 2001
From: Juan Martinez Ramirez <juan.martinezramirez@snowflake.com>
Date: Tue, 22 Oct 2024 17:30:53 -0600
Subject: [PATCH] Added property to LoginRequestData to specify httpRequest
 timeout.

---
 Snowflake.Data/Core/Authenticator/IAuthenticator.cs |  2 +-
 .../Core/Authenticator/MFACacheAuthenticator.cs     |  2 ++
 Snowflake.Data/Core/RestRequest.cs                  |  3 +++
 Snowflake.Data/Core/Session/SFSession.cs            | 13 +++++++++++++
 4 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/Snowflake.Data/Core/Authenticator/IAuthenticator.cs b/Snowflake.Data/Core/Authenticator/IAuthenticator.cs
index 0de637b35..f6497e980 100644
--- a/Snowflake.Data/Core/Authenticator/IAuthenticator.cs
+++ b/Snowflake.Data/Core/Authenticator/IAuthenticator.cs
@@ -145,7 +145,7 @@ private SFRestRequest BuildLoginRequest()
             };
             SetSpecializedAuthenticatorData(ref data);
 
-            return session.BuildTimeoutRestRequest(loginUrl, new LoginRequest() { data = data });
+            return data.HttpTimeout.HasValue ? session.BuildTimeoutRestRequest(loginUrl, new LoginRequest() { data = data }, data.HttpTimeout.Value) : session.BuildTimeoutRestRequest(loginUrl, new LoginRequest() { data = data });
         }
     }
 
diff --git a/Snowflake.Data/Core/Authenticator/MFACacheAuthenticator.cs b/Snowflake.Data/Core/Authenticator/MFACacheAuthenticator.cs
index d4b679632..f72509b59 100644
--- a/Snowflake.Data/Core/Authenticator/MFACacheAuthenticator.cs
+++ b/Snowflake.Data/Core/Authenticator/MFACacheAuthenticator.cs
@@ -2,6 +2,7 @@
  * Copyright (c) 2024 Snowflake Computing Inc. All rights reserved.
  */
 
+using System;
 using System.Threading;
 using System.Threading.Tasks;
 using Snowflake.Data.Core.Tools;
@@ -34,6 +35,7 @@ protected override void SetSpecializedAuthenticatorData(ref LoginRequestData dat
             // Only need to add the password to Data for basic authentication
             data.password = session.properties[SFSessionProperty.PASSWORD];
             data.SessionParameters[SFSessionParameter.CLIENT_REQUEST_MFA_TOKEN] = true;
+            data.HttpTimeout = TimeSpan.FromSeconds(60);
             if (!string.IsNullOrEmpty(session._mfaToken?.ToString()))
             {
                 data.Token = SecureStringHelper.Decode(session._mfaToken);
diff --git a/Snowflake.Data/Core/RestRequest.cs b/Snowflake.Data/Core/RestRequest.cs
index b26feae43..de988895b 100644
--- a/Snowflake.Data/Core/RestRequest.cs
+++ b/Snowflake.Data/Core/RestRequest.cs
@@ -268,6 +268,9 @@ class LoginRequestData
         [JsonProperty(PropertyName = "SESSION_PARAMETERS", NullValueHandling = NullValueHandling.Ignore)]
         internal Dictionary<SFSessionParameter, Object> SessionParameters { get; set; }
 
+        [JsonIgnore]
+        internal TimeSpan? HttpTimeout { get; set; }
+
         public override string ToString()
         {
             return String.Format("LoginRequestData {{ClientAppVersion: {0},\n AccountName: {1},\n loginName: {2},\n ClientEnv: {3},\n authenticator: {4} }}",
diff --git a/Snowflake.Data/Core/Session/SFSession.cs b/Snowflake.Data/Core/Session/SFSession.cs
index 5a1c8f7f0..f09e6cd2f 100644
--- a/Snowflake.Data/Core/Session/SFSession.cs
+++ b/Snowflake.Data/Core/Session/SFSession.cs
@@ -460,6 +460,19 @@ internal SFRestRequest BuildTimeoutRestRequest(Uri uri, Object body)
             };
         }
 
+        internal SFRestRequest BuildTimeoutRestRequest(Uri uri, Object body, TimeSpan httpTimeout)
+        {
+            return new SFRestRequest()
+            {
+                jsonBody = body,
+                Url = uri,
+                authorizationToken = SF_AUTHORIZATION_BASIC,
+                RestTimeout = connectionTimeout,
+                HttpTimeout = httpTimeout,
+                _isLogin = true
+            };
+        }
+
         internal void UpdateSessionParameterMap(List<NameValueParameter> parameterList)
         {
             logger.Debug("Update parameter map");