Add Attach Database functionality to Object Management Service (#2193)

src/Microsoft.SqlTools.ServiceLayer/Admin/AdminService.cs

@@ -8,10 +8,13 @@
 using System;
 using System.Collections.Concurrent;
 using System.Threading.Tasks;
+using Microsoft.Data.SqlClient;
+using Microsoft.SqlServer.Management.Common;
 using Microsoft.SqlServer.Management.Smo;
 using Microsoft.SqlTools.Hosting.Protocol;
 using Microsoft.SqlTools.ServiceLayer.Admin.Contracts;
 using Microsoft.SqlTools.ServiceLayer.Connection;
+using Microsoft.SqlTools.ServiceLayer.DisasterRecovery;
 using Microsoft.SqlTools.ServiceLayer.Hosting;
 using Microsoft.SqlTools.ServiceLayer.Management;
 using Microsoft.SqlTools.ServiceLayer.Utility;
@@ -71,6 +74,8 @@ namespace Microsoft.SqlTools.ServiceLayer.Admin
             serviceHost.SetRequestHandler(CreateLoginRequest.Type, HandleCreateLoginRequest, true);
             serviceHost.SetRequestHandler(DefaultDatabaseInfoRequest.Type, HandleDefaultDatabaseInfoRequest, true);
             serviceHost.SetRequestHandler(GetDatabaseInfoRequest.Type, HandleGetDatabaseInfoRequest, true);
+            serviceHost.SetRequestHandler(GetDataFolderRequest.Type, HandleGetDataFolderRequest, true);
+            serviceHost.SetRequestHandler(GetAssociatedFilesRequest.Type, HandleGetAssociatedFilesRequest, true);
         }
 
         /// <summary>
@@ -153,6 +158,74 @@ namespace Microsoft.SqlTools.ServiceLayer.Admin
             });
         }
 
+        /// <summary>
+        /// Handle get database info request
+        /// </summary>
+        internal static async Task HandleGetDataFolderRequest(
+            GetDataFolderParams databaseParams,
+            RequestContext<string> requestContext)
+        {
+            Func<Task> requestHandler = async () =>
+            {
+                ConnectionInfo connInfo;
+                AdminService.ConnectionServiceInstance.TryFindConnection(
+                        databaseParams.ConnectionUri,
+                        out connInfo);
+                using (SqlConnection sqlConn = ConnectionService.OpenSqlConnection(connInfo))
+                {
+                    // Connection gets disconnected when backup is done
+                    ServerConnection serverConnection = new ServerConnection(sqlConn);
+                    var dataFolder = CommonUtilities.GetDefaultDataFolder(serverConnection);
+                    await requestContext.SendResult(dataFolder);
+                }
+            };
+
+            Task task = Task.Run(async () => await requestHandler()).ContinueWithOnFaulted(async t =>
+            {
+                // Get innermost exception to get original error message
+                Exception ex = t.Exception;
+                while (ex.InnerException != null)
+                {
+                    ex = ex.InnerException;
+                };
+                await requestContext.SendError(ex.Message);
+            });
+        }
+
+        /// <summary>
+        /// Handle get associated database files request
+        /// </summary>
+        internal static async Task HandleGetAssociatedFilesRequest(
+            GetAssociatedFilesParams databaseParams,
+            RequestContext<string[]> requestContext)
+        {
+            Func<Task> requestHandler = async () =>
+            {
+                ConnectionInfo connInfo;
+                AdminService.ConnectionServiceInstance.TryFindConnection(
+                        databaseParams.ConnectionUri,
+                        out connInfo);
+                using (SqlConnection sqlConn = ConnectionService.OpenSqlConnection(connInfo))
+                {
+                    // Connection gets disconnected when backup is done
+                    ServerConnection serverConnection = new ServerConnection(sqlConn);
+                    var files = CommonUtilities.GetAssociatedFilePaths(serverConnection, databaseParams.PrimaryFilePath);
+                    await requestContext.SendResult(files);
+                }
+            };
+
+            Task task = Task.Run(async () => await requestHandler()).ContinueWithOnFaulted(async t =>
+            {
+                // Get innermost exception to get original error message
+                Exception ex = t.Exception;
+                while (ex.InnerException != null)
+                {
+                    ex = ex.InnerException;
+                };
+                await requestContext.SendError(ex.Message);
+            });
+        }
+
         /// <summary>
         /// Return database info for a specific database
         /// </summary>

src/Microsoft.SqlTools.ServiceLayer/Admin/Contracts/GetAssociatedFilesRequest.cs

@@ -0,0 +1,36 @@
+//
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+//
+
+#nullable disable
+
+using Microsoft.SqlTools.Hosting.Protocol.Contracts;
+
+namespace Microsoft.SqlTools.ServiceLayer.Admin.Contracts
+{
+    /// <summary>
+    /// Params for a get database info request
+    /// </summar>
+    public class GetAssociatedFilesParams
+    {
+        /// <summary>
+        /// URI identifier for the connection to the target server.
+        /// </summary>
+        public string ConnectionUri { get; set; }
+        /// <summary>
+        /// The file path for the primary file that we want to get the associated files for.
+        /// </summary>
+        public string PrimaryFilePath { get; set; }
+    }
+
+    /// <summary>
+    /// Get database info request mapping
+    /// </summary>
+    public class GetAssociatedFilesRequest
+    {
+        public static readonly
+            RequestType<GetAssociatedFilesParams, string[]> Type =
+                RequestType<GetAssociatedFilesParams, string[]>.Create("admin/getassociatedfiles");
+    }
+}

src/Microsoft.SqlTools.ServiceLayer/Admin/Contracts/GetDataFolderRequest.cs

@@ -0,0 +1,32 @@
+//
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+//
+
+#nullable disable
+
+using Microsoft.SqlTools.Hosting.Protocol.Contracts;
+
+namespace Microsoft.SqlTools.ServiceLayer.Admin.Contracts
+{
+    /// <summary>
+    /// Params for a get database info request
+    /// </summar>
+    public class GetDataFolderParams
+    {
+        /// <summary>
+        /// URI identifier for the connection to get the server folder info for
+        /// </summary>
+        public string ConnectionUri { get; set; }
+    }
+
+    /// <summary>
+    /// Get database info request mapping
+    /// </summary>
+    public class GetDataFolderRequest
+    {
+        public static readonly
+            RequestType<GetDataFolderParams, string> Type =
+                RequestType<GetDataFolderParams, string>.Create("admin/getdatafolder");
+    }
+}

src/Microsoft.SqlTools.ServiceLayer/DisasterRecovery/CommonUtilities.cs

@@ -14,6 +14,7 @@ using Microsoft.SqlServer.Management.Sdk.Sfc;
 using Microsoft.SqlServer.Management.Smo;
 using SMO = Microsoft.SqlServer.Management.Smo;
 using Microsoft.SqlTools.ServiceLayer.Management;
+using System.IO;
 
 namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery
 {
@@ -101,7 +102,7 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery
             }
             else
             {
-                return (RestoreItemLocation+RestoreItemDeviceType.ToString() + IsLogicalDevice.ToString()).GetHashCode();
+                return (RestoreItemLocation + RestoreItemDeviceType.ToString() + IsLogicalDevice.ToString()).GetHashCode();
             }
         }
     }
@@ -226,7 +227,7 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery
                 }
                 return false;
             }
-            catch(Exception)
+            catch (Exception)
             {
                 return false;
             }
@@ -241,7 +242,7 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery
                 DataSet ds = new DataSet();
                 ds.Locale = System.Globalization.CultureInfo.InvariantCulture;
                 req.Urn = "Server/BackupDevice";
-                ds = en.Process(this.sqlConnection,req);
+                ds = en.Process(this.sqlConnection, req);
 
                 if (ds.Tables[0].Rows.Count > 0)
                 {
@@ -249,7 +250,7 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery
                 }
                 return false;
             }
-            catch(Exception)
+            catch (Exception)
             {
                 return false;
             }
@@ -342,6 +343,69 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery
             return backupFolder;
         }
 
+        public static string GetDefaultDataFolder(ServerConnection connection)
+        {
+            string dataFolder = string.Empty;
+
+            Enumerator en = null;
+            DataSet ds = new DataSet();
+            ds.Locale = System.Globalization.CultureInfo.InvariantCulture;
+            Request req = new Request();
+            en = new Enumerator();
+            req.Urn = "Server/Setting";
+            ds = en.Process(connection, req);
+
+            if (ds.Tables[0].Rows.Count > 0)
+            {
+                dataFolder = Convert.ToString(ds.Tables[0].Rows[0]["DefaultFile"], System.Globalization.CultureInfo.InvariantCulture);
+            }
+            return dataFolder;
+        }
+
+        public static string[] GetAssociatedFilePaths(ServerConnection connection, string primaryFilePath)
+        {
+            var databaseFiles = new List<string>();
+
+            var enumerator = new Enumerator();
+            var req = new Request()
+            {
+                Urn = $"Server/PrimaryFile[@Name='{Urn.EscapeString(primaryFilePath)}']/File",
+                Fields = new string[] { "IsFile", "FileName" }
+            };
+            var dataTable = (DataTable)enumerator.Process(connection, req);
+
+            foreach (DataRow currentRow in dataTable.Rows)
+            {
+                var primaryFolder = Path.GetDirectoryName(primaryFilePath);
+                var originalPath = (string)currentRow["FileName"];
+                var originalFileName = Path.GetFileName(originalPath);
+                var filePath = Path.Join(primaryFolder, originalFileName);
+
+                // Check if file exists with the constructed path.
+                // If it's an XI (XStore Integration) path, then assume it exists, otherwise retrieve info for the file to check if it exists.
+                var exists = true;
+                var isXIPath = PathWrapper.IsXIPath(primaryFilePath);
+                if (!isXIPath)
+                {
+                    var request = new Request()
+                    {
+                        Urn = string.Format(System.Globalization.CultureInfo.CurrentCulture, "Server/File[@FullName='{0}']", Urn.EscapeString(filePath)),
+                        Fields = new string[] { "IsFile" }
+                    };
+
+                    DataTable data = (new Enumerator()).Process(connection, request);
+
+                    // If the enumerator could find the file, then it exists
+                    exists = data?.Rows.Count > 0;
+                }
+                if (exists)
+                {
+                    databaseFiles.Add(filePath);
+                }
+            }
+            return databaseFiles.ToArray();
+        }
+
         public int GetMediaRetentionValue()
         {
             int afterDays = 0;
@@ -354,7 +418,7 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery
 
                 req.Urn = "Server/Configuration";
                 ds = en.Process(this.sqlConnection, req);
-                for (int i = 0 ; i < ds.Tables[0].Rows.Count; i++)
+                for (int i = 0; i < ds.Tables[0].Rows.Count; i++)
                 {
                     if (Convert.ToString(ds.Tables[0].Rows[i]["Name"], System.Globalization.CultureInfo.InvariantCulture) == "media retention")
                     {
@@ -694,11 +758,14 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery
             {
                 switch (type)
                 {
-                    case "D": fileType = RestoreConstants.Data;
+                    case "D":
+                        fileType = RestoreConstants.Data;
                         break;
-                    case "S": fileType = RestoreConstants.FileStream;
+                    case "S":
+                        fileType = RestoreConstants.FileStream;
                         break;
-                    default: fileType = RestoreConstants.NotKnown;
+                    default:
+                        fileType = RestoreConstants.NotKnown;
                         break;
                 }
             }
@@ -707,7 +774,7 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery
         public BackupsetType GetBackupsetTypeFromBackupTypesOnDevice(int type)
         {
             BackupsetType Result = BackupsetType.BackupsetDatabase;
-            switch(type)
+            switch (type)
             {
                 case 1:
                     Result = BackupsetType.BackupsetDatabase;
@@ -732,7 +799,7 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery
         public BackupsetType GetBackupsetTypeFromBackupTypesOnHistory(string type)
         {
             BackupsetType result = BackupsetType.BackupsetDatabase;
-            switch(type)
+            switch (type)
             {
                 case "D":
                     result = BackupsetType.BackupsetDatabase;
@@ -760,7 +827,7 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery
             DataSet backupsetfiles = new DataSet();
             backupsetfiles.Locale = System.Globalization.CultureInfo.InvariantCulture;
 
-            if(backupsetId > 0)
+            if (backupsetId > 0)
             {
                 req.Urn = "Server/BackupSet[@ID='" + Urn.EscapeString(Convert.ToString(backupsetId, System.Globalization.CultureInfo.InvariantCulture)) + "']/File";
             }
@@ -796,7 +863,7 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery
 
             ArrayList sources = new ArrayList();
             DataSet backupSet = GetBackupSetById(backupsetId);
-            if(backupSet.Tables[0].Rows.Count == 1)
+            if (backupSet.Tables[0].Rows.Count == 1)
             {
                 string mediaSetID = Convert.ToString(backupSet.Tables[0].Rows[0]["MediaSetId"], System.Globalization.CultureInfo.InvariantCulture);
 
@@ -805,12 +872,12 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery
                 DataSet mediafamily = new DataSet();
                 mediafamily.Locale = System.Globalization.CultureInfo.InvariantCulture;
 
-                req.Urn = "Server/BackupMediaSet[@ID='"+Urn.EscapeString(mediaSetID)+"']/MediaFamily";
+                req.Urn = "Server/BackupMediaSet[@ID='" + Urn.EscapeString(mediaSetID) + "']/MediaFamily";
                 mediafamily = en.Process(this.sqlConnection, req);
 
                 if (mediafamily.Tables[0].Rows.Count > 0)
                 {
-                    for (int j = 0 ; j < mediafamily.Tables[0].Rows.Count; j ++)
+                    for (int j = 0; j < mediafamily.Tables[0].Rows.Count; j++)
                     {
                         RestoreItemSource itemSource = new RestoreItemSource();
                         itemSource.RestoreItemLocation = Convert.ToString(mediafamily.Tables[0].Rows[j]["PhysicalDeviceName"], System.Globalization.CultureInfo.InvariantCulture);
@@ -870,7 +937,7 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery
             backupSets.Locale = System.Globalization.CultureInfo.InvariantCulture;
             OrderBy orderByBackupDate;
 
-            req.Urn = "Server/BackupSet[@Name='"+Urn.EscapeString(backupSetName)+"' and @DatabaseName='"+ Urn.EscapeString(databaseName)+"']";
+            req.Urn = "Server/BackupSet[@Name='" + Urn.EscapeString(backupSetName) + "' and @DatabaseName='" + Urn.EscapeString(databaseName) + "']";
             req.OrderByList = new OrderBy[1];
             orderByBackupDate = new OrderBy("BackupFinishDate", OrderBy.Direction.Desc);
             req.OrderByList[0] = orderByBackupDate;
@@ -957,7 +1024,7 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery
             ds.Locale = System.Globalization.CultureInfo.InvariantCulture;
             Request req = new Request();
 
-            req.Urn = "Server/Login[@Name='"+Urn.EscapeString(loginName)+"']";
+            req.Urn = "Server/Login[@Name='" + Urn.EscapeString(loginName) + "']";
             req.Fields = new string[1];
             req.Fields[0] = "DefaultDatabase";
             ds = en.Process(this.sqlConnection, req);
@@ -998,7 +1065,7 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery
             Request req = new Request();
             ArrayList result = null;
             int count = 0;
-            req.Urn = "Server/BackupDevice[@PhysicalLocation='" +Urn.EscapeString(physicalPath)+ "']";
+            req.Urn = "Server/BackupDevice[@PhysicalLocation='" + Urn.EscapeString(physicalPath) + "']";
 
             ds = en.Process(this.sqlConnection, req);
             count = ds.Tables[0].Rows.Count;
@@ -1006,7 +1073,7 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery
             if (count > 0)
             {
                 result = new ArrayList(count);
-                for(int i = 0; i < count; i++)
+                for (int i = 0; i < count; i++)
                 {
                     result.Add(Convert.ToString(ds.Tables[0].Rows[0]["Name"], System.Globalization.CultureInfo.InvariantCulture));
                 }

src/Microsoft.SqlTools.ServiceLayer/ObjectManagement/Contracts/AttachDatabaseRequest.cs

@@ -0,0 +1,29 @@
+//
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+//
+
+#nullable disable
+using Microsoft.SqlTools.Hosting.Protocol.Contracts;
+
+namespace Microsoft.SqlTools.ServiceLayer.ObjectManagement.Contracts
+{
+    public class DatabaseFileData
+    {
+        public string DatabaseName { get; set; }
+        public string[] DatabaseFilePaths { get; set; }
+        public string Owner { get; set; }
+    }
+
+    public class AttachDatabaseRequestParams
+    {
+        public string ConnectionUri { get; set; }
+        public DatabaseFileData[] Databases { get; set; }
+        public bool GenerateScript { get; set; }
+    }
+
+    public class AttachDatabaseRequest
+    {
+        public static readonly RequestType<AttachDatabaseRequestParams, string> Type = RequestType<AttachDatabaseRequestParams, string>.Create("objectManagement/attachDatabase");
+    }
+}

src/Microsoft.SqlTools.ServiceLayer/ObjectManagement/ObjectManagementService.cs

@@ -70,6 +70,7 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectManagement
             this.serviceHost.SetRequestHandler(DisposeViewRequest.Type, HandleDisposeViewRequest, true);
             this.serviceHost.SetRequestHandler(SearchRequest.Type, HandleSearchRequest, true);
             this.serviceHost.SetRequestHandler(DetachDatabaseRequest.Type, HandleDetachDatabaseRequest, true);
+            this.serviceHost.SetRequestHandler(AttachDatabaseRequest.Type, HandleAttachDatabaseRequest, true);
             this.serviceHost.SetRequestHandler(DropDatabaseRequest.Type, HandleDropDatabaseRequest, true);
         }
 
@@ -207,6 +208,13 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectManagement
             await requestContext.SendResult(sqlScript);
         }
 
+        internal async Task HandleAttachDatabaseRequest(AttachDatabaseRequestParams requestParams, RequestContext<string> requestContext)
+        {
+            var handler = this.GetObjectTypeHandler(SqlObjectType.Database) as DatabaseHandler;
+            var sqlScript = handler.Attach(requestParams);
+            await requestContext.SendResult(sqlScript);
+        }
+
         internal async Task HandleDropDatabaseRequest(DropDatabaseRequestParams requestParams, RequestContext<string> requestContext)
         {
             var handler = this.GetObjectTypeHandler(SqlObjectType.Database) as DatabaseHandler;

src/Microsoft.SqlTools.ServiceLayer/ObjectManagement/ObjectTypes/Database/DatabaseHandler.cs

@@ -20,6 +20,7 @@ using Microsoft.SqlTools.Utility;
 using System.Text;
 using System.IO;
 using Microsoft.SqlTools.ServiceLayer.Utility.SqlScriptFormatters;
+using System.Collections.Specialized;
 using Microsoft.SqlTools.SqlCore.Utility;
 
 namespace Microsoft.SqlTools.ServiceLayer.ObjectManagement
@@ -393,6 +394,60 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectManagement
             return builder.ToString();
         }
 
+        public string Attach(AttachDatabaseRequestParams attachParams)
+        {
+            var sqlScript = string.Empty;
+            ConnectionInfo connectionInfo = this.GetConnectionInfo(attachParams.ConnectionUri);
+            using (var dataContainer = CreateDatabaseDataContainer(attachParams.ConnectionUri, null, true, null))
+            {
+                var server = dataContainer.Server!;
+                var originalExecuteMode = server.ConnectionContext.SqlExecutionModes;
+                if (attachParams.GenerateScript)
+                {
+                    server.ConnectionContext.SqlExecutionModes = SqlExecutionModes.CaptureSql;
+                    server.ConnectionContext.CapturedSql.Clear();
+                }
+                try
+                {
+                    foreach (var database in attachParams.Databases)
+                    {
+                        var fileCollection = new StringCollection();
+                        fileCollection.AddRange(database.DatabaseFilePaths);
+                        if (database.Owner != SR.general_default)
+                        {
+                            server.AttachDatabase(database.DatabaseName, fileCollection, database.Owner);
+                        }
+                        else
+                        {
+                            server.AttachDatabase(database.DatabaseName, fileCollection);
+                        }
+                    }
+                    if (attachParams.GenerateScript)
+                    {
+                        var builder = new StringBuilder();
+                        var capturedText = server.ConnectionContext.CapturedSql.Text;
+                        foreach (var entry in capturedText)
+                        {
+                            if (entry != null)
+                            {
+                                builder.AppendLine(entry);
+                            }
+                        }
+                        sqlScript = builder.ToString();
+                    }
+                }
+                finally
+                {
+                    if (attachParams.GenerateScript)
+                    {
+                        server.ConnectionContext.SqlExecutionModes = originalExecuteMode;
+                    }
+                    dataContainer.ServerConnection.Disconnect();
+                }
+            }
+            return sqlScript;
+        }
+
         /// <summary>
         /// Used to drop the specified database
         /// </summary>

test/Microsoft.SqlTools.ServiceLayer.IntegrationTests/ObjectManagement/DatabaseHandlerTests.cs

@@ -3,9 +3,12 @@
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 //
 
+#nullable disable
+
 using System.Collections.Generic;
 using System.Collections.Specialized;
 using System.Linq;
+using System.Text;
 using System.Threading.Tasks;
 using Microsoft.Data.SqlClient;
 using Microsoft.SqlServer.Management.Common;
@@ -500,6 +503,107 @@ namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.ObjectManagement
             }
         }
 
+        [Test]
+        [TestCase(true)]
+        [TestCase(false)]
+        public async Task AttachDatabaseTest(bool generateScript)
+        {
+            using (SqlTestDb testDb = await SqlTestDb.CreateNewAsync(TestServerType.OnPrem, false, null, null, nameof(AttachDatabaseTest)))
+            {
+                var connectionResult = await LiveConnectionHelper.InitLiveConnectionInfoAsync("master", serverType: TestServerType.OnPrem);
+                using (SqlConnection sqlConn = ConnectionService.OpenSqlConnection(connectionResult.ConnectionInfo))
+                {
+                    var serverConn = new ServerConnection(sqlConn);
+                    var server = new Server(serverConn);
+                    var objUrn = ObjectManagementTestUtils.GetDatabaseURN(testDb.DatabaseName);
+                    var database = server.GetSmoObject(objUrn) as Database;
+
+                    var originalOwner = database!.Owner;
+                    var originalFilePaths = new List<string>();
+                    foreach (FileGroup group in database.FileGroups)
+                    {
+                        foreach (DataFile file in group.Files)
+                        {
+                            originalFilePaths.Add(file.FileName);
+                        }
+                    }
+                    foreach (LogFile file in database.LogFiles)
+                    {
+                        originalFilePaths.Add(file.FileName);
+                    }
+
+                    // Detach database so that we can re-attach it with the database handler method.
+                    // Have to set database to single user mode to close active connections before detaching it.
+                    database.DatabaseOptions.UserAccess = SqlServer.Management.Smo.DatabaseUserAccess.Single;
+                    database.Alter(TerminationClause.RollbackTransactionsImmediately);
+                    server.DetachDatabase(testDb.DatabaseName, false);
+                    var dbExists = this.DatabaseExists(testDb.DatabaseName, server);
+                    Assert.That(dbExists, Is.False, "Database was not correctly detached before doing attach test.");
+
+                    try
+                    {
+                        var handler = new DatabaseHandler(ConnectionService.Instance);
+                        var attachParams = new AttachDatabaseRequestParams()
+                        {
+                            ConnectionUri = connectionResult.ConnectionInfo.OwnerUri,
+                            Databases = new DatabaseFileData[]
+                            {
+                                new DatabaseFileData()
+                                {
+                                    Owner = originalOwner,
+                                    DatabaseName = testDb.DatabaseName,
+                                    DatabaseFilePaths = originalFilePaths.ToArray()
+                                }
+                            },
+                            GenerateScript = generateScript
+                        };
+                        var script = handler.Attach(attachParams);
+
+                        if (generateScript)
+                        {
+                            dbExists = this.DatabaseExists(testDb.DatabaseName, server);
+                            Assert.That(dbExists, Is.False, "Should not have attached DB when only generating a script.");
+
+                            var queryBuilder = new StringBuilder(); 
+                            queryBuilder.AppendLine("USE [master]");
+                            queryBuilder.AppendLine($"CREATE DATABASE [{testDb.DatabaseName}] ON ");
+
+                            for (int i = 0; i < originalFilePaths.Count - 1; i++)
+                            {
+                                var file = originalFilePaths[i];
+                                queryBuilder.AppendLine($"( FILENAME = N'{file}' ),");
+                            }
+                            queryBuilder.AppendLine($"( FILENAME = N'{originalFilePaths[originalFilePaths.Count - 1]}' )");
+
+                            queryBuilder.AppendLine(" FOR ATTACH");
+                            queryBuilder.AppendLine($"if exists (select name from master.sys.databases sd where name = N'{testDb.DatabaseName}' and SUSER_SNAME(sd.owner_sid) = SUSER_SNAME() ) EXEC [{testDb.DatabaseName}].dbo.sp_changedbowner @loginame=N'{originalOwner}', @map=false");
+
+                            Assert.That(script, Is.EqualTo(queryBuilder.ToString()), "Did not get expected attach database script");
+                        }
+                        else
+                        {
+                            Assert.That(script, Is.Empty, "Should not have generated a script for this Attach operation.");
+
+                            server.Databases.Refresh();
+                            dbExists = this.DatabaseExists(testDb.DatabaseName, server);
+                            Assert.That(dbExists, "Database was not attached successfully");
+                        }
+                    }
+                    finally
+                    {
+                        dbExists = this.DatabaseExists(testDb.DatabaseName, server);
+                        if (!dbExists)
+                        {
+                            // Reattach database so it can get dropped during cleanup
+                            var fileCollection = new StringCollection();
+                            originalFilePaths.ForEach(file => fileCollection.Add(file));
+                            server.AttachDatabase(testDb.DatabaseName, fileCollection);
+                        }
+                    }
+                }
+            }
+        }
+
         [Test]
         public async Task DeleteDatabaseTest()
         {

test/Microsoft.SqlTools.ServiceLayer.IntegrationTests/ObjectManagement/UtilsTests.cs

@@ -0,0 +1,135 @@
+//
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+//
+
+#nullable disable
+
+using System.Collections.Generic;
+using System.Collections.Specialized;
+using System.IO;
+using System.Threading.Tasks;
+using Microsoft.Data.SqlClient;
+using Microsoft.SqlServer.Management.Common;
+using Microsoft.SqlServer.Management.Smo;
+using Microsoft.SqlTools.ServiceLayer.Connection;
+using Microsoft.SqlTools.ServiceLayer.DisasterRecovery;
+using Microsoft.SqlTools.ServiceLayer.IntegrationTests.Utility;
+using Microsoft.SqlTools.ServiceLayer.Test.Common;
+using NUnit.Framework;
+
+namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.ObjectManagement
+{
+    public class UtilsTests
+    {
+        [Test]
+        public async Task GetDataFolderTest()
+        {
+            using (SqlTestDb testDb = await SqlTestDb.CreateNewAsync(TestServerType.OnPrem, false, null, null, nameof(GetDataFolderTest)))
+            {
+                var connectionResult = await LiveConnectionHelper.InitLiveConnectionInfoAsync(testDb.DatabaseName, serverType: TestServerType.OnPrem);
+                using (SqlConnection sqlConn = ConnectionService.OpenSqlConnection(connectionResult.ConnectionInfo))
+                {
+                    var serverConn = new ServerConnection(sqlConn);
+                    var server = new Server(serverConn);
+                    var objUrn = ObjectManagementTestUtils.GetDatabaseURN(testDb.DatabaseName);
+                    var database = server.GetSmoObject(objUrn) as Database;
+
+                    var dataFilePath = database.FileGroups[0].Files[0].FileName;
+                    var expectedDataFolder = Path.GetDirectoryName(dataFilePath).ToString();
+
+                    var actualDataFolder = CommonUtilities.GetDefaultDataFolder(serverConn);
+                    actualDataFolder = Path.TrimEndingDirectorySeparator(actualDataFolder);
+                    Assert.That(actualDataFolder, Is.EqualTo(expectedDataFolder).IgnoreCase, "Did not get expected data file folder path.");
+                }
+            }
+        }
+
+        [Test]
+        public async Task GetAssociatedFilesTest()
+        {
+            using (SqlTestDb testDb = await SqlTestDb.CreateNewAsync(TestServerType.OnPrem, false, null, null, nameof(GetAssociatedFilesTest)))
+            {
+                var connectionResult = await LiveConnectionHelper.InitLiveConnectionInfoAsync(testDb.DatabaseName, serverType: TestServerType.OnPrem);
+                using (SqlConnection sqlConn = ConnectionService.OpenSqlConnection(connectionResult.ConnectionInfo))
+                {
+                    var serverConn = new ServerConnection(sqlConn);
+                    var server = new Server(serverConn);
+                    var objUrn = ObjectManagementTestUtils.GetDatabaseURN(testDb.DatabaseName);
+                    var database = server.GetSmoObject(objUrn) as Database;
+
+                    var expectedFilePaths = new List<string>();
+                    DataFile primaryFile = null;
+                    foreach (FileGroup group in database.FileGroups)
+                    {
+                        foreach (DataFile file in group.Files)
+                        {
+                            expectedFilePaths.Add(file.FileName);
+                            if (file.IsPrimaryFile)
+                            {
+                                primaryFile = file;
+                            }
+                        }
+                    }
+                    foreach (LogFile file in database.LogFiles)
+                    {
+                        expectedFilePaths.Add(file.FileName);
+                    }
+
+                    // Detach database so that we don't throw an error when trying to access the primary data file
+                    // Have to set database to single user mode to close active connections before detaching it.
+                    database.DatabaseOptions.UserAccess = SqlServer.Management.Smo.DatabaseUserAccess.Single;
+                    database.Alter(TerminationClause.RollbackTransactionsImmediately);
+                    server.DetachDatabase(testDb.DatabaseName, false);
+                    try
+                    {
+                        Assert.That(primaryFile, Is.Not.Null, "Could not find a primary file in the list of database files.");
+                        var actualFilePaths = CommonUtilities.GetAssociatedFilePaths(serverConn, primaryFile.FileName);
+                        Assert.That(actualFilePaths, Is.EqualTo(expectedFilePaths).IgnoreCase, "The list of associated files did not match the actual files for the database.");
+                    }
+                    finally
+                    {
+                        // Reattach database so it can get dropped during cleanup
+                        var fileCollection = new StringCollection();
+                        expectedFilePaths.ForEach(file => fileCollection.Add(file));
+                        server.AttachDatabase(testDb.DatabaseName, fileCollection);
+                    }
+                }
+            }
+        }
+
+        [Test]
+        public async Task ThrowErrorWhenDatabaseExistsTest()
+        {
+            using (SqlTestDb testDb = await SqlTestDb.CreateNewAsync(TestServerType.OnPrem, false, null, null, nameof(ThrowErrorWhenDatabaseExistsTest)))
+            {
+                var connectionResult = await LiveConnectionHelper.InitLiveConnectionInfoAsync(testDb.DatabaseName, serverType: TestServerType.OnPrem);
+                using (SqlConnection sqlConn = ConnectionService.OpenSqlConnection(connectionResult.ConnectionInfo))
+                {
+                    var serverConn = new ServerConnection(sqlConn);
+                    var server = new Server(serverConn);
+                    var objUrn = ObjectManagementTestUtils.GetDatabaseURN(testDb.DatabaseName);
+                    var database = server.GetSmoObject(objUrn) as Database;
+
+                    DataFile primaryFile = null;
+                    foreach (FileGroup group in database.FileGroups)
+                    {
+                        foreach (DataFile file in group.Files)
+                        {
+                            if (file.IsPrimaryFile)
+                            {
+                                primaryFile = file;
+                            }
+                        }
+                    }
+
+                    Assert.That(
+                        () => CommonUtilities.GetAssociatedFilePaths(serverConn, primaryFile.FileName),
+                        Throws.Exception,
+                        "Should throw an error when trying to open a database file that's already in use."
+                    );
+                }
+            }
+        }
+    }
+}