From f0007af7a14f1e33e61dfb1fc4ffe6a3c57d942c Mon Sep 17 00:00:00 2001
From: Juan Martinez Ramirez <juan.martinezramirez@snowflake.com>
Date: Fri, 19 Jul 2024 16:19:47 -0600
Subject: [PATCH] Initial support for loading token from token_file_path for
 oauth authentication

---
 .../Core/SnowflakeTomlConnectionBuilder.cs    | 31 +++++++++++++++++++
 1 file changed, 31 insertions(+)

diff --git a/Snowflake.Data/Core/SnowflakeTomlConnectionBuilder.cs b/Snowflake.Data/Core/SnowflakeTomlConnectionBuilder.cs
index 9d702be47..020c4f0a3 100644
--- a/Snowflake.Data/Core/SnowflakeTomlConnectionBuilder.cs
+++ b/Snowflake.Data/Core/SnowflakeTomlConnectionBuilder.cs
@@ -7,6 +7,7 @@ namespace Snowflake.Data.Core
     using System;
     using System.Collections.Generic;
     using System.IO;
+    using System.Linq;
     using System.Text;
     using Tomlyn;
     using Tomlyn.Model;
@@ -16,12 +17,15 @@ public class SnowflakeTomlConnectionBuilder
     {
         private const string DefaultConnectionName = "default";
         private const string DefaultSnowflakeFolder = ".snowflake";
+        private const string DefaultTokenPath = "/snowflake/session/token";
 
         private Dictionary<string, string> TomlToNetPropertiesMapper = new Dictionary<string, string>(StringComparer.InvariantCultureIgnoreCase)
         {
             { "DATABASE", "DB" }
         };
 
+
+
         private readonly FileOperations _fileOperations;
         private readonly EnvironmentOperations _environmentOperations;
 
@@ -50,15 +54,42 @@ public string GetConnectionStringFromToml(string connectionName = null)
         private string GetConnectionStringFromTomlTable(TomlTable connectionToml)
         {
             var connectionStringBuilder = new StringBuilder();
+            var tokenFilePathValue = string.Empty;
+            var isOauth = connectionToml.TryGetValue("authenticator", out var authenticator) && authenticator.ToString().Equals("oauth");
             foreach (var property in connectionToml.Keys)
             {
+                if (isOauth && property.Equals("token_file_path", StringComparison.InvariantCultureIgnoreCase))
+                {
+                    tokenFilePathValue = (string)connectionToml[property];
+                    continue;
+                }
                 var mappedProperty = TomlToNetPropertiesMapper.TryGetValue(property, out var mapped) ? mapped : property;
                 connectionStringBuilder.Append($"{mappedProperty}={(string)connectionToml[property]};");
             }
 
+            if (!isOauth || connectionToml.ContainsKey("token"))
+                return connectionStringBuilder.ToString();
+
+            var token = LoadTokenFromFile(tokenFilePathValue);
+            if (!string.IsNullOrEmpty(token))
+            {
+                connectionStringBuilder.Append($"token={token};");
+            }
+            else
+            {
+                // log warning TODO
+            }
+
+
             return connectionStringBuilder.ToString();
         }
 
+        private string LoadTokenFromFile(string tokenFilePathValue)
+        {
+            var tokenFile = _fileOperations.Exists(tokenFilePathValue) ? tokenFilePathValue : DefaultTokenPath;
+            return _fileOperations.Exists(tokenFile) ? _fileOperations.ReadAllText(tokenFile) : null;
+        }
+
         private TomlTable GetTomlTableFromConfig(string tomlPath, string connectionName)
         {
             TomlTable result = null;