Hotfix for ads dashboard for sqlondemand (#885)

* Hotfix for ads dashboard for sqlondemand

* Generalizing engine edition functions and adapting tests

* Minor comment and cosmetic fixes

* Adding SqlOnDemand server edition string

src/Microsoft.SqlTools.ManagedBatchParser/ReliableConnection/CachedServerInfo.cs

@@ -9,6 +9,8 @@ using System.Data;
 using Microsoft.Data.SqlClient;
 using System.Diagnostics;
 using System.Linq;
+using Microsoft.SqlServer.Management.Common;
+using Microsoft.SqlServer.Management.Dmf;
 using Microsoft.SqlTools.Utility;
 
 namespace Microsoft.SqlTools.ServiceLayer.Connection.ReliableConnection
@@ -37,7 +39,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection.ReliableConnection
         }
 
         public enum CacheVariable {
-            IsSqlDw,
+            EngineEdition,
             IsAzure,
             IsCloud
         }
@@ -101,7 +103,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection.ReliableConnection
         {
             public bool IsAzure;
             public DateTime LastUpdate;
-            public bool IsSqlDw;
+            public DatabaseEngineEdition EngineEdition;
         }
 
         private const int _maxCacheSize = 1024;
@@ -163,27 +165,51 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection.ReliableConnection
             AddOrUpdateCache(connection, isAzure, CacheVariable.IsAzure);
         }
 
-        public void AddOrUpdateIsSqlDw(IDbConnection connection, bool isSqlDw)
+        public void AddOrUpdateEngineEdition(IDbConnection connection, DatabaseEngineEdition engineEdition)
         {
-            AddOrUpdateCache(connection, isSqlDw, CacheVariable.IsSqlDw);
+            AddOrUpdateCache(connection, engineEdition, CacheVariable.EngineEdition);
         }
 
-        private void AddOrUpdateCache(IDbConnection connection, bool newState, CacheVariable cacheVar)
+        private void AddOrUpdateCache(IDbConnection connection, object newState, CacheVariable cacheVar)
         {
             Validate.IsNotNull(nameof(connection), connection);
             SqlConnectionStringBuilder builder = new SqlConnectionStringBuilder(connection.ConnectionString);
             AddOrUpdateCache(builder, newState, cacheVar);
         }
 
-        public void AddOrUpdateCache(SqlConnectionStringBuilder builder, bool newState, CacheVariable cacheVar)
+        private bool IsAppropriateType(object newState, CacheVariable cacheVar)
         {
+            if (newState is DatabaseEngineEdition)
+            {
+                return cacheVar == CacheVariable.EngineEdition;
+            }
+
+            if (newState is bool)
+            {
+                return cacheVar == CacheVariable.IsAzure || cacheVar == CacheVariable.IsCloud;
+            }
+
+            return false;
+        }
+
+        private T ConvertState<T>(object state)
+        {
+            return (T) Convert.ChangeType(state, typeof(T));
+        }
+
+        public void AddOrUpdateCache(SqlConnectionStringBuilder builder, object newState, CacheVariable cacheVar)
+        {
+            if (!IsAppropriateType(newState, cacheVar))
+            {
+                throw new FunctionWrongArgumentTypeException("AddOrUpdateCache: mismatch between expected type of CacheVariable and the type of provided update object");
+            }
             Validate.IsNotNull(nameof(builder), builder);
             Validate.IsNotNullOrWhitespaceString(nameof(builder) + ".DataSource", builder.DataSource);
             CachedInfo info;
             bool hasFound = TryGetCacheValue(builder, out info);
 
-            if ((cacheVar == CacheVariable.IsSqlDw && hasFound && info.IsSqlDw == newState) || 
-                (cacheVar == CacheVariable.IsAzure && hasFound && info.IsAzure == newState)) 
+            if (cacheVar == CacheVariable.EngineEdition && hasFound && info.EngineEdition == ConvertState<DatabaseEngineEdition>(newState) ||
+                cacheVar == CacheVariable.IsAzure && hasFound && info.IsAzure == ConvertState<bool>(newState))
             {
                 // No change needed
                 return;
@@ -196,13 +222,13 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection.ReliableConnection
                     CacheKey key = new CacheKey(builder);
                     CleanupCache(key);
 
-                    if (cacheVar == CacheVariable.IsSqlDw)
+                    if (cacheVar == CacheVariable.EngineEdition)
                     {
-                        info.IsSqlDw = newState;
+                        info.EngineEdition = ConvertState<DatabaseEngineEdition>(newState);
                     }
                     else if (cacheVar == CacheVariable.IsAzure)
                     {
-                        info.IsAzure = newState;
+                        info.IsAzure = ConvertState<bool>(newState);
                     }
                     info.LastUpdate = DateTime.UtcNow;
                     _cache.AddOrUpdate(key, info, (k, oldValue) => info);
@@ -232,15 +258,15 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection.ReliableConnection
             }
         }
 
-        public bool TryGetIsSqlDw(IDbConnection connection, out bool isSqlDw)
+        public DatabaseEngineEdition TryGetEngineEdition(IDbConnection connection, out DatabaseEngineEdition engineEdition)
         {
             Validate.IsNotNull(nameof(connection), connection);
 
             SqlConnectionStringBuilder builder = new SqlConnectionStringBuilder(connection.ConnectionString);
-            return TryGetIsSqlDw(builder, out isSqlDw);
+            return TryGetEngineEdition(builder, out engineEdition);
         }
 
-        public bool TryGetIsSqlDw(SqlConnectionStringBuilder builder, out bool isSqlDw)
+        public DatabaseEngineEdition TryGetEngineEdition(SqlConnectionStringBuilder builder, out DatabaseEngineEdition engineEdition)
         {
             Validate.IsNotNull(nameof(builder), builder);
             Validate.IsNotNullOrWhitespaceString(nameof(builder) + ".DataSource", builder.DataSource);
@@ -249,12 +275,11 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection.ReliableConnection
 
             if(hasFound)
             {
-                isSqlDw = info.IsSqlDw;
-                return true;
+                engineEdition = info.EngineEdition;
+                return engineEdition;
             }
-            
-            isSqlDw = false;
-            return false;
+
+            return engineEdition = DatabaseEngineEdition.Unknown;
         }
 
         private static SqlConnectionStringBuilder SafeGetConnectionStringFromConnection(IDbConnection connection)

src/Microsoft.SqlTools.ManagedBatchParser/ReliableConnection/ReliableConnectionHelper.cs

@@ -421,6 +421,42 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection.ReliableConnection
             cmd.CommandTimeout = CachedServerInfo.Instance.GetQueryTimeoutSeconds(cmd.Connection);
         }
 
+        public static DatabaseEngineEdition GetEngineEdition(IDbConnection connection)
+        {
+            Validate.IsNotNull(nameof(connection), connection);
+            if (!(connection.State == ConnectionState.Open))
+            {
+                Logger.Write(TraceEventType.Warning, Resources.ConnectionPassedToIsCloudShouldBeOpen);
+            }
+
+            Func<string, DatabaseEngineEdition> executeCommand = commandText =>
+            {
+                DatabaseEngineEdition result = DatabaseEngineEdition.Unknown;
+                ExecuteReader(connection,
+                    commandText,
+                    readResult: (reader) =>
+                    {
+                        reader.Read();
+                        result = (DatabaseEngineEdition)int.Parse(reader[0].ToString(), CultureInfo.InvariantCulture);
+                    }
+                );
+                return result;
+            };
+
+            DatabaseEngineEdition engineEdition = DatabaseEngineEdition.Unknown;
+            try
+            {
+                engineEdition = executeCommand(SqlConnectionHelperScripts.EngineEdition);
+            }
+            catch (SqlException)
+            {
+                // The default query contains a WITH (NOLOCK).  This doesn't work for Azure DW or SqlOnDemand, so when things don't work out, 
+                // we'll fall back to a version without NOLOCK and try again.
+                engineEdition = executeCommand(SqlConnectionHelperScripts.EngineEditionWithLock);
+            }
+
+            return engineEdition;
+        }
 
         /// <summary>
         /// Return true if the database is an Azure database
@@ -429,98 +465,13 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection.ReliableConnection
         /// <returns></returns>
         public static bool IsCloud(IDbConnection connection)
         {
-            Validate.IsNotNull(nameof(connection), connection);
-            if (!(connection.State == ConnectionState.Open))
-            {
-                Logger.Write(TraceEventType.Warning, Resources.ConnectionPassedToIsCloudShouldBeOpen);
-            }
-
-            Func<string, bool> executeCommand = commandText =>
-            {
-                bool result = false;
-                ExecuteReader(connection,
-                          commandText,
-                          readResult: (reader) =>
-                          {
-                              reader.Read();
-                              int engineEditionId = int.Parse(reader[0].ToString(), CultureInfo.InvariantCulture);
-
-                              result = IsCloudEngineId(engineEditionId);
-                          }
-                );
-                return result;
-            };
-
-            bool isSqlCloud = false;
-            try
-            {
-                isSqlCloud = executeCommand(SqlConnectionHelperScripts.EngineEdition);
-            }
-            catch (SqlException)
-            {
-                // The default query contains a WITH (NOLOCK).  This doesn't work for Azure DW, so when things don't work out, 
-                // we'll fall back to a version without NOLOCK and try again.
-                isSqlCloud = executeCommand(SqlConnectionHelperScripts.EngineEditionWithLock);
-            }
-
-            return isSqlCloud;
+            return IsCloudEngineId((int)GetEngineEdition(connection));
         }
-
         private static bool IsCloudEngineId(int engineEditionId)
         {
             return cloudEditions.Value.Contains(engineEditionId);
         }
 
-        /// <summary>
-        /// Determines if the type of database that a connection is being made to is SQL data warehouse.
-        /// </summary>
-        /// <param name="connection"></param>
-        /// <returns>True if the database is a SQL data warehouse</returns>
-        public static bool IsSqlDwDatabase(IDbConnection connection)
-        {
-            Validate.IsNotNull(nameof(connection), connection);
-
-            Func<string, bool> executeCommand = commandText =>
-            {
-                bool result = false;
-                ExecuteReader(connection,
-                              commandText,
-                              readResult: (reader) =>
-                              {
-                                  reader.Read();
-                                  int engineEditionId = int.Parse(reader[0].ToString(), CultureInfo.InvariantCulture);
-
-                                  result = IsSqlDwEngineId(engineEditionId);
-                              }
-                    );
-                return result;
-            };
-
-            bool isSqlDw = false;
-            try
-            {
-                isSqlDw = executeCommand(SqlConnectionHelperScripts.EngineEdition);
-            }
-            catch (SqlException)
-            {
-                // The default query contains a WITH (NOLOCK).  This doesn't work for Azure DW, so when things don't work out, 
-                // we'll fall back to a version without NOLOCK and try again.
-                isSqlDw = executeCommand(SqlConnectionHelperScripts.EngineEditionWithLock);
-            }
-
-            return isSqlDw;
-        }
-
-        /// <summary>
-        /// Compares the engine edition id of a given database with that of SQL data warehouse.
-        /// </summary>
-        /// <param name="engineEditionId"></param>
-        /// <returns>True if the engine edition id is that of SQL data warehouse</returns>
-        private static bool IsSqlDwEngineId(int engineEditionId)
-        {
-            return engineEditionId == SqlDwEngineEditionId;
-        }
-
         /// <summary>
         /// Handles the exceptions typically thrown when a SQLConnection is being opened
         /// </summary>

src/Microsoft.SqlTools.ManagedBatchParser/ReliableConnection/ReliableSqlConnection.cs

@@ -33,6 +33,7 @@ using System.Globalization;
 using System.Text;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.SqlServer.Management.Common;
 using Microsoft.SqlTools.Utility;
 
 namespace Microsoft.SqlTools.ServiceLayer.Connection.ReliableConnection
@@ -49,7 +50,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection.ReliableConnection
         private readonly RetryPolicy _connectionRetryPolicy;
         private RetryPolicy _commandRetryPolicy;
         private Guid _azureSessionId;
-        private bool _isSqlDwDatabase;
+        private DatabaseEngineEdition _engineEdition;
 
         /// <summary>
         /// Initializes a new instance of the ReliableSqlConnection class with a given connection string
@@ -104,23 +105,33 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection.ReliableConnection
             }
         }
 
+        private DatabaseEngineEdition GetOrReadCachedEngineEdition(IDbConnection conn)
+        {
+            if (CachedServerInfo.Instance.TryGetEngineEdition(conn, out _engineEdition) == DatabaseEngineEdition.Unknown)
+            {
+                _engineEdition = ReliableConnectionHelper.GetEngineEdition(conn);
+                CachedServerInfo.Instance.AddOrUpdateEngineEdition(conn, _engineEdition);
+            }
+
+            return _engineEdition;
+        }
+
         /// <summary>
         /// Determines if a connection is being made to a SQL DW database.
         /// </summary>
         /// <param name="conn">A connection object.</param>
         private bool IsSqlDwConnection(IDbConnection conn)
         {
-            //Set the connection only if it has not been set earlier.
-            //This is assuming that it is highly unlikely for a connection to change between instances.
-            //Hence any subsequent calls to this method will just return the cached value and not 
-            //verify again if this is a SQL DW database connection or not.
-            if (!CachedServerInfo.Instance.TryGetIsSqlDw(conn, out _isSqlDwDatabase))
-            {
-                _isSqlDwDatabase = ReliableConnectionHelper.IsSqlDwDatabase(conn);
-                CachedServerInfo.Instance.AddOrUpdateIsSqlDw(conn, _isSqlDwDatabase);;
-            }
+            return GetOrReadCachedEngineEdition(conn) == DatabaseEngineEdition.SqlDataWarehouse;
+        }
 
-            return _isSqlDwDatabase;
+        /// <summary>
+        /// Determines if a connection is being made to SQLOnDemand.
+        /// </summary>
+        /// <param name="conn">A connection object.</param>
+        private bool IsSqlOnDemandConnection(IDbConnection conn)
+        {
+            return GetOrReadCachedEngineEdition(conn) == DatabaseEngineEdition.SqlOnDemand;
         }
 
         [System.Diagnostics.CodeAnalysis.SuppressMessage("Microsoft.Security", "CA2100:Review SQL queries for security vulnerabilities")]
@@ -430,7 +441,7 @@ SET NUMERIC_ROUNDABORT OFF;";
                 using (IDbCommand command = CreateReliableCommand())
                 {
                     IDbConnection connection = command.Connection;
-                    if (!IsSqlDwConnection(connection))
+                    if (!IsSqlDwConnection(connection) && !IsSqlOnDemandConnection(connection))
                     {
                         command.CommandText = QueryAzureSessionId;
                         object result = command.ExecuteScalar();
@@ -524,7 +535,7 @@ SET NUMERIC_ROUNDABORT OFF;";
             Tuple<string,bool>[] sessionSettings = new Tuple<string,bool>[2];
 
             IDbConnection connection = originalCommand.Connection;
-            if (IsSqlDwConnection(connection))
+            if (IsSqlDwConnection(connection) || IsSqlOnDemandConnection(connection))
             {
                 // SESSIONPROPERTY is not supported. Use default values for now
                 sessionSettings[0] = Tuple.Create("ANSI_NULLS", true);