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.CoreServices/Connection/ConnectionServiceCore.cs

@@ -477,9 +477,12 @@ namespace Microsoft.SqlTools.CoreServices.Connection
                     case (int) DatabaseEngineEdition.SqlDataWarehouse:
                         serverEdition = SR.AzureSqlDwEdition;
                         break;
-                    case (int) DatabaseEngineEdition.SqlStretchDatabase:
+                    case (int)DatabaseEngineEdition.SqlStretchDatabase:
                         serverEdition = SR.AzureSqlStretchEdition;
                         break;
+                    case (int)DatabaseEngineEdition.SqlOnDemand:
+                        serverEdition = SR.AzureSqlAnalyticsOnDemandEdition;
+                        break;
                     default:
                         serverEdition =  SR.AzureSqlDbEdition;
                         break;

src/Microsoft.SqlTools.CoreServices/Localization/sr.cs

@@ -101,6 +101,14 @@ namespace Microsoft.SqlTools.CoreServices
             }
         }
 
+        public static string AzureSqlAnalyticsOnDemandEdition
+        {
+            get
+            {
+                return Keys.GetString(Keys.AzureSqlAnalyticsOnDemandEdition);
+            }
+        }
+
         public static string QueryServiceCancelAlreadyCompleted
         {
             get
@@ -3764,6 +3772,9 @@ namespace Microsoft.SqlTools.CoreServices
             public const string AzureSqlStretchEdition = "AzureSqlStretchEdition";
 
 
+            public const string AzureSqlAnalyticsOnDemandEdition = "AzureSqlAnalyticsOnDemandEdition";
+
+
             public const string QueryServiceCancelAlreadyCompleted = "QueryServiceCancelAlreadyCompleted";
 
 

src/Microsoft.SqlTools.CoreServices/Localization/sr.strings

@@ -51,6 +51,7 @@ ConnectionParamsValidateNullSqlAuth(string component) = {0} cannot be null or em
 AzureSqlDbEdition = Azure SQL DB
 AzureSqlDwEdition = Azure SQL Data Warehouse
 AzureSqlStretchEdition = Azure SQL Stretch Database
+AzureSqlAnalyticsOnDemandEdition = Azure SQL Analytics on-demand
 
 ############################################################################
 # Query Execution Service

src/Microsoft.SqlTools.CoreServices/Localization/sr.xlf

@@ -2296,6 +2296,11 @@
         <target state="new">Azure SQL Stretch Database</target>
         <note></note>
       </trans-unit>
+      <trans-unit id="AzureSqlAnalyticsOnDemandEdition">
+        <source>Azure SQL Analytics on-demand</source>
+        <target state="new">Azure SQL Analytics on-demand</target>
+        <note></note>
+      </trans-unit>
       <trans-unit id="Error_InvalidDirectoryName">
         <source>Path {0} is not a valid directory </source>
         <target state="new">Path {0} is not a valid directory </target>

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) || 
+            if (cacheVar == CacheVariable.EngineEdition && hasFound && info.EngineEdition == ConvertState<DatabaseEngineEdition>(newState) ||
-                (cacheVar == CacheVariable.IsAzure && hasFound && info.IsAzure == 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;
+                engineEdition = info.EngineEdition;
-                return true;
+                return engineEdition;
             }
-            
+
-            isSqlDw = false;
+            return engineEdition = DatabaseEngineEdition.Unknown;
-            return false;
         }
 
         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);
+            return IsCloudEngineId((int)GetEngineEdition(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;
         }
-
         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.
+            return GetOrReadCachedEngineEdition(conn) == DatabaseEngineEdition.SqlDataWarehouse;
-            //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 _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);

src/Microsoft.SqlTools.ServiceLayer/Admin/Database/DatabasePrototype.cs

@@ -2278,10 +2278,13 @@ WHERE do.database_id = @DbID
             {
                 try
                 {
-                    foreach (LogFile logfile in database.LogFiles)
+                    if (this.context.Server.DatabaseEngineEdition != DatabaseEngineEdition.SqlOnDemand)
                     {
-                        DatabaseFilePrototype logfilePrototype = new DatabaseFilePrototype(this, logfile);
+                        foreach (LogFile logfile in database.LogFiles)
-                        this.Add(logfilePrototype);
+                        {
+                            DatabaseFilePrototype logfilePrototype = new DatabaseFilePrototype(this, logfile);
+                            this.Add(logfilePrototype);
+                        }
                     }
                 }
                 catch (ExecutionFailureException)

src/Microsoft.SqlTools.ServiceLayer/Connection/ConnectionService.cs

@@ -473,6 +473,9 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
                     case (int)DatabaseEngineEdition.SqlStretchDatabase:
                         serverEdition = SR.AzureSqlStretchEdition;
                         break;
+                    case (int)DatabaseEngineEdition.SqlOnDemand:
+                        serverEdition = SR.AzureSqlAnalyticsOnDemandEdition;
+                        break;
                     default:
                         serverEdition = SR.AzureSqlDbEdition;
                         break;

src/Microsoft.SqlTools.ServiceLayer/Localization/sr.cs

@@ -99,6 +99,14 @@ namespace Microsoft.SqlTools.ServiceLayer
             {
                 return Keys.GetString(Keys.AzureSqlStretchEdition);
             }
+        }
+
+        public static string AzureSqlAnalyticsOnDemandEdition
+        {
+            get
+            {
+                return Keys.GetString(Keys.AzureSqlAnalyticsOnDemandEdition);
+            }
         }
 
         public static string QueryServiceCancelAlreadyCompleted
@@ -3161,7 +3169,10 @@ namespace Microsoft.SqlTools.ServiceLayer
             public const string AzureSqlDwEdition = "AzureSqlDwEdition";
 
 
-            public const string AzureSqlStretchEdition = "AzureSqlStretchEdition";
+            public const string AzureSqlStretchEdition = "AzureSqlStretchEdition";
+
+
+            public const string AzureSqlAnalyticsOnDemandEdition = "AzureSqlAnalyticsOnDemandEdition";
 
 
             public const string QueryServiceCancelAlreadyCompleted = "QueryServiceCancelAlreadyCompleted";

src/Microsoft.SqlTools.ServiceLayer/Localization/sr.strings

@@ -51,6 +51,7 @@ ConnectionParamsValidateNullSqlAuth(string component) = {0} cannot be null or em
 AzureSqlDbEdition = Azure SQL DB
 AzureSqlDwEdition = Azure SQL Data Warehouse
 AzureSqlStretchEdition = Azure SQL Stretch Database
+AzureSqlAnalyticsOnDemandEdition = Azure SQL Analytics on-demand
 
 ############################################################################
 # Query Execution Service

src/Microsoft.SqlTools.ServiceLayer/Localization/sr.xlf

@@ -1590,6 +1590,11 @@
         <source>Azure SQL Stretch Database</source>
         <target state="new">Azure SQL Stretch Database</target>
         <note></note>
+      </trans-unit>	  
+      <trans-unit id="AzureSqlAnalyticsOnDemandEdition">
+        <source>Azure SQL Analytics on-demand</source>
+        <target state="new">Azure SQL Analytics on-demand</target>
+        <note></note>
       </trans-unit>
       <trans-unit id="EditDataValueTooLarge">
         <source>Value {0} is too large to fit in column of type {1}</source>

test/Microsoft.SqlTools.ServiceLayer.UnitTests/Connection/CachedServerInfoTests.cs

@@ -6,6 +6,7 @@
 using Microsoft.SqlTools.ServiceLayer.Connection.ReliableConnection;
 using Xunit;
 using Microsoft.Data.SqlClient;
+using Microsoft.SqlServer.Management.Common;
 
 namespace Microsoft.SqlTools.ServiceLayer.UnitTests.Connection
 {
@@ -26,45 +27,43 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.Connection
         {
             // Set sqlDw result into cache
             string dataSource = "testDataSource";
-            bool isSqlDwResult;
+            DatabaseEngineEdition engineEdition;
             SqlConnectionStringBuilder testSource = new SqlConnectionStringBuilder
             {
                 DataSource = dataSource,
                 InitialCatalog = string.Empty
             };
-            cache.AddOrUpdateCache(testSource, true, CachedServerInfo.CacheVariable.IsSqlDw);
+            cache.AddOrUpdateCache(testSource, DatabaseEngineEdition.SqlDataWarehouse, CachedServerInfo.CacheVariable.EngineEdition);
 
             // Expect the same returned result
-            Assert.True(cache.TryGetIsSqlDw(testSource, out isSqlDwResult));
+            Assert.Equal(cache.TryGetEngineEdition(testSource, out engineEdition), DatabaseEngineEdition.SqlDataWarehouse);
-            Assert.True(isSqlDwResult);
 
             // And expect the same for the null string
-            Assert.True(cache.TryGetIsSqlDw(new SqlConnectionStringBuilder
+            Assert.Equal(cache.TryGetEngineEdition(new SqlConnectionStringBuilder
             {
                 DataSource = dataSource
                 // Initial Catalog is null. Can't set explicitly as this throws
-            }, out isSqlDwResult));
+            }, out engineEdition), DatabaseEngineEdition.SqlDataWarehouse);
-            Assert.True(isSqlDwResult);
 
-            // But expect false for a different DB
+            // But expect NotEqual for a different DB
-            Assert.False(cache.TryGetIsSqlDw(new SqlConnectionStringBuilder
+            Assert.NotEqual(cache.TryGetEngineEdition(new SqlConnectionStringBuilder
             {
                 DataSource = dataSource,
                 InitialCatalog = "OtherDb"
-            }, out isSqlDwResult));
+            }, out engineEdition), DatabaseEngineEdition.SqlDataWarehouse);
         }
 
         [Theory]
-        [InlineData(null, true)]  // is SqlDW instance
+        [InlineData(null, DatabaseEngineEdition.SqlDataWarehouse)]  // is SqlDW instance
-        [InlineData("", true)]  // is SqlDW instance
+        [InlineData("", DatabaseEngineEdition.SqlDataWarehouse)]  // is SqlDW instance
-        [InlineData("myDb", true)]  // is SqlDW instance
+        [InlineData("myDb", DatabaseEngineEdition.SqlDataWarehouse)]  // is SqlDW instance
-        [InlineData(null, false)]   // is not a SqlDw Instance
+        [InlineData(null, DatabaseEngineEdition.SqlOnDemand)]   // is SqlOnDemand Instance
-        [InlineData("", false)]   // is not a SqlDw Instance
+        [InlineData("", DatabaseEngineEdition.SqlOnDemand)]   // is SqlOnDemand Instance
-        [InlineData("myDb", false)]  // is not SqlDW instance
+        [InlineData("myDb", DatabaseEngineEdition.SqlOnDemand)]  // is SqlOnDemand instance
-        public void AddOrUpdateIsSqlDw(string dbName, bool state)
+        public void AddOrUpdateEngineEditiopn(string dbName, DatabaseEngineEdition state)
         {
-            // Set sqlDw result into cache
+            // Set result into cache
-            bool isSqlDwResult;
+            DatabaseEngineEdition engineEdition;
             SqlConnectionStringBuilder testSource = new SqlConnectionStringBuilder
             {
                 DataSource = "testDataSource"
@@ -74,41 +73,41 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.Connection
                 testSource.InitialCatalog = dbName;
             }
 
-            cache.AddOrUpdateCache(testSource, state, CachedServerInfo.CacheVariable.IsSqlDw);
+            cache.AddOrUpdateCache(testSource, state, CachedServerInfo.CacheVariable.EngineEdition);
 
             // Expect the same returned result
-            Assert.True(cache.TryGetIsSqlDw(testSource, out isSqlDwResult));
+            Assert.NotEqual(cache.TryGetEngineEdition(testSource, out engineEdition), DatabaseEngineEdition.Unknown);
-            Assert.Equal(isSqlDwResult, state);
+            Assert.Equal(engineEdition, state);
         }
 
         [Theory]
-        [InlineData(true)]  // is SqlDW instance
+        [InlineData(DatabaseEngineEdition.SqlDataWarehouse)]  // is SqlDW instance
-        [InlineData(false)]   // is not a SqlDw Instance
+        [InlineData(DatabaseEngineEdition.SqlOnDemand)]   // is SqlOnDemand Instance
-        public void AddOrUpdateIsSqlDwFalseToggle(bool state)
+        public void AddOrUpdateEngineEditionToggle(DatabaseEngineEdition state)
         {
-            // Set sqlDw result into cache
+            // Set result into cache
-            bool isSqlDwResult;
+            DatabaseEngineEdition engineEdition;
             SqlConnectionStringBuilder testSource = new SqlConnectionStringBuilder
             {
                 DataSource = "testDataSource"
             };
-            cache.AddOrUpdateCache(testSource, state, CachedServerInfo.CacheVariable.IsSqlDw);
+            cache.AddOrUpdateCache(testSource, state, CachedServerInfo.CacheVariable.EngineEdition);
 
             // Expect the same returned result
-            Assert.True(cache.TryGetIsSqlDw(testSource, out isSqlDwResult));
+            Assert.NotEqual(cache.TryGetEngineEdition(testSource, out engineEdition), DatabaseEngineEdition.Unknown);
-            Assert.Equal(isSqlDwResult, state);
+            Assert.Equal(engineEdition, state);
 
-            // Toggle isSqlDw cache state
+            DatabaseEngineEdition newState = state == DatabaseEngineEdition.SqlDataWarehouse ?
-            bool isSqlDwResultToggle;
+                DatabaseEngineEdition.SqlOnDemand : DatabaseEngineEdition.SqlDataWarehouse;
-            cache.AddOrUpdateCache(testSource, !state, CachedServerInfo.CacheVariable.IsSqlDw);
 
-            // Expect the oppisite returned result
+            cache.AddOrUpdateCache(testSource, newState, CachedServerInfo.CacheVariable.EngineEdition);
-            Assert.True(cache.TryGetIsSqlDw(testSource, out isSqlDwResultToggle));
-            Assert.Equal(isSqlDwResultToggle, !state);
 
+            // Expect the opposite returned result
+            Assert.NotEqual(cache.TryGetEngineEdition(testSource, out engineEdition), DatabaseEngineEdition.Unknown);
+            Assert.Equal(engineEdition, newState);
         }
 
-       // [Fact]
+       /* [Fact]
         public void AddOrUpdateIsSqlDwFalseToggle()
         {
             bool state = true;
@@ -148,16 +147,17 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.Connection
             Assert.Equal(isSqlDwResult3, !state);
 
         }
+        */
 
         [Fact]
-        public void AskforSqlDwBeforeCached()
+        public void AskforEngineEditionBeforeCached()
         {
-            bool isSqlDwResult;
+            DatabaseEngineEdition engineEdition;
-            Assert.False(cache.TryGetIsSqlDw(new SqlConnectionStringBuilder
+            Assert.Equal(cache.TryGetEngineEdition(new SqlConnectionStringBuilder
             {
                 DataSource = "testDataSourceUnCached"
             }, 
-            out isSqlDwResult));
+            out engineEdition), DatabaseEngineEdition.Unknown);
         }
     }
 }