Restore Database to AvailabilityGroups

#Require -Version 5
#Require -Modules SqlPS, FailoverClusters, ActiveDirectory

function Deploy-Database{
    param(
        [Parameter(Mandatory=$true)]
        [string]$DestServer,
        [Parameter(Mandatory=$true)]
        [string]$SourceServer,
        [Parameter(Mandatory=$true)]
        [string]$SourceDatabase,
        [ValidateSet(2050,2060,2070,1433)]
        [int]$Port = 1433,
        [string]$BackupLocation = "\\NetworkLocation\Backups"
    )
    begin{

    }
    process{

        [string]$backupFile = "$BackupLocation\$SourceDatabase.bak"

        if([System.IO.File]::Exists($backupFile)){
            [System.IO.File]::Delete($backupFile);
        }

        #ensure source is pingable
        if(!(Test-Connection -ComputerName $SourceServer -Count 2 -Quiet)){
            throw "Source server $SourceServer could not be contacted.  Please verify that this server is correct or is reachable."
        }

        #ensure destination is pingable
        if(!(Test-Connection -ComputerName $DestServer -Count 2 -Quiet)){
            throw "Destination server $DestServer could not be contacted.  Please verify that this server is correct or is reachable."
        }

        #ensure db exist on source and is online and can be backed up
        $srv = New-Object Microsoft.SqlServer.Management.Smo.Server "$SourceServer,$Port"
        $db = $srv.Databases | ?{$_.Name -eq $SourceDatabase -and $_.Status -eq [Microsoft.SqlServer.Management.Smo.DatabaseStatus]::Normal}
        if(!$db){
            throw "Database $SourceDatabase on server $SourceServer does not exist, or is not in a state that allows backups."
        }
        
        #backup source database
        $srv.ConnectionContext.StatementTimeout = 0;
        Backup-SqlDatabase -InputObject $srv -Database $SourceDatabase -BackupFile $backupFile -CompressionOption On 
        $srv.ConnectionContext.Disconnect();

        $srv = New-Object Microsoft.SqlServer.Management.smo.Server "$DestServer,$Port"
        $cluster = gwmi -class "MSCluster_Cluster" -namespace "root\mscluster" -computername $DestServer -Authentication PacketPrivacy -ErrorAction SilentlyContinue

        if($cluster -and $srv.IsHadrEnabled){
            $sqlNodes = New-Object System.Collections.Generic.List[object]
            $agGroups = New-Object System.Collections.Generic.List[object]

            #if multiple ag groups, restore to the ag group with the least number of databases
            foreach($ag in $srv.AvailabilityGroups){
                [void]$agGroups.Add([PSCustomObject]@{
                    Name = $ag.Name 
                    PrimaryName = $ag.PrimaryReplicaServerName
                    DatabaseCount = $ag.AvalabilityDatabases.Count 
                });
            }
            $agGroup = $agGroups | sort DatabaseCount | select -First 1

            $srv.ConnectionContext.Disconnect();
            
            #get nodes in cluster to restore the db onto
            Get-Cluster -Name $cluster.Name | Get-ClusterNode | %{
                $nodeName = $_.Name 
                [void]$sqlNodes.Add([PSCustomObject]@{
                    ServerName = $nodeName
                });
            }

            #test the sql connectivity on the nodes before trying any restores. If any fail, boot.
            #additionally, test to make sure that the database doesn't already exist on any of the nodes.
            $sqlNodes | %{
                $serverName = $_.ServerName 
                $srvConn = New-Object Microsoft.SqlServer.Management.Common.ServerConnection "$serverName,$Port"
                try{
                    $srvConn.Connect();
                    $srv = New-Object Microsoft.SqlServer.Management.Smo.Server $srvConn
                    if($srv.Databases[$SourceDatabase]){
                        throw "Database $SourceDatabase already exists on server $serverName.  Restore will not continue."
                    }
                }
                catch{
                    throw "Could not connect to server $ServerName.  Database will not be restored.  Exception is:`r`n$($_ | fl -force | out-string)"
                }
                finally{
                    $srvConn.Disconnect();
                }
            }

            #restore with norecovery to all nodes
            $sqlNodes | %{
                $serverName = $_.ServerName 
                $srv = New-Object Microsoft.SqlServer.Management.Smo.Server "$serverName,$Port"
                $srv.ConnectionContext.StatementTimeout = 0;
                Restore-SqlDatabase -InputObject $srv -Database $SourceDatabase -BackupFile $backupFile -NoRecovery
                $srv.ConnectionContext.Disconnect();
            }

            #connect to primary of group and recover just that database
            $srv = New-Object Microsoft.SqlServer.Management.Smo.Server "$($agGroup.PrimaryName),$Port"
            $db = $srv.Databases[$SourceDatabase]  #just another check...
            if(!$db){
                throw "Database $SourceDatabase on server $($agGroup.PrimaryName) does not exist.  If the database $SourceDatabase exists on servers $($sqlNodes -join ',') they will need to be dropped before another deploy can be attempted."
            }
            $srv.ConnectionContext.ExecuteNonQuery("RESTORE DATABASE [$SourceDatabase] WITH RECOVERY;");
            $srv.ConnectionContext.Disconnect();

            #add primary to AG
            Add-SqlAvailabilityDatabase -Database $SourceDatabase -Path "SQLSERVER:\SQL\$($agGroup.PrimaryName)`,$Port\DEFAULT\AvailabilityGroups\$($agGroup.Name)" 

            #add subsequent nodes to AG
            $sqlNodes | ?{$_.ServerName -ne $agGroup.PrimaryName} | %{
                $serverName = $_.ServerName 
                Add-SqlAvailabilityDatabase -Database $SourceDatabase -Path "SQLSERVER:\SQL\$ServerName`,$Port\DEFAULT\AvailabilityGroups\$($agGroup.Name)"  
            }


        }
        else{
            #just a single instance restore....
            $srv = New-Object Microsoft.SqlServer.Management.Smo.Server "$DestServer,$Port"
            $srv.ConnectionContext.StatementTimeout = 0;
            if($srv.Databases[$SourceDatabase]){
                throw "Database $SourceDatabase already exists on server $DestServer."
            }
            Restore-SqlDatabase -InputObject $srv -Database $SourceDatabase -BackupFile $backupFile
            $srv.ConnectionContext.Disconnect();
        }

    }
    end{

    }
}