Dba Query Wrapper

Invoke-SqlAsync.ps1

function Invoke-SqlAsync {
    param (
        [Alias('Connection', 'Conn')]
        [ValidateNotNullOrEmpty()]
        [System.Data.SqlClient.SQLConnection]$SQLConnection,

        [Parameter(Mandatory = $true, Position = 0, ParameterSetName = "Query")]
        [string]
        $Query,

        [ValidateSet("DataSet", "DataTable", "DataRow", "PSObject", "SingleValue")]
        [string]
        $As = "DataRow",

        [System.Collections.IDictionary]
        $SqlParameters,

        [switch]
        $AppendServerInstance,

        [Int32]$QueryTimeout = 600,


        [switch]$MessagesToOutput,

        [Alias('Silent')]
        [switch]
        $EnableException
    )

    begin {

        if ($As -eq "PSObject") {
            #This code scrubs DBNulls.  Props to Dave Wyatt
            $cSharp = @'
                using System;
                using System.Data;
                using System.Management.Automation;

                public class DBNullScrubber
                {
                    public static PSObject DataRowToPSObject(DataRow row)
                    {
                        PSObject psObject = new PSObject();

                        if (row != null && (row.RowState & DataRowState.Detached) != DataRowState.Detached)
                        {
                            foreach (DataColumn column in row.Table.Columns)
                            {
                                Object value = null;
                                if (!row.IsNull(column))
                                {
                                    value = row[column];
                                }

                                psObject.Properties.Add(new PSNoteProperty(column.ColumnName, value));
                            }
                        }

                        return psObject;
                    }
                }
'@

            try {
                Add-Type -TypeDefinition $cSharp -ReferencedAssemblies 'System.Data', 'System.Xml' -ErrorAction stop
            }
            catch {
                if (-not $_.ToString() -like "*The type name 'DBNullScrubber' already exists*") {
                    Write-Warning "Could not load DBNullScrubber.  Defaulting to DataRow output: $_."
                    $As = "Datarow"
                }
            }
        }

        #Handle existing connections

        $GoSplitterRegex = [regex]'(?smi)^[\s]*GO[\s]*$'

    }
    process {
        $Conn = $SQLConnection
        if ($ParseGO) {
            Write-Message -Level Debug -Message "Stripping GOs from source"
            $Pieces = $GoSplitterRegex.Split($Query)
        }
        else {
            $Pieces = , $Query
        }
        # Only execute non-empty statements
        $Pieces = $Pieces | Where-Object { $_.Trim().Length -gt 0 }
        foreach ($piece in $Pieces) {
            $cmd = New-Object system.Data.SqlClient.SqlCommand($piece, $Conn)
            $cmd.CommandTimeout = $QueryTimeout

            if ($null -ne $SqlParameters) {
                $SqlParameters.GetEnumerator() |
                    ForEach-Object {
                    if ($null -ne $_.Value) {
                        $cmd.Parameters.AddWithValue($_.Key, $_.Value)
                    }
                    else {
                        $cmd.Parameters.AddWithValue($_.Key, [DBNull]::Value)
                    }
                } > $null
            }

            $ds = New-Object system.Data.DataSet
            $da = New-Object system.Data.SqlClient.SqlDataAdapter($cmd)

            try {
                Write-Message -Level VeryVerbose -Message "start invoke query"
                if ($MessagesToOutput) {
                    $pool = [RunspaceFactory]::CreateRunspacePool(1, [int]$env:NUMBER_OF_PROCESSORS + 1)
                    $pool.ApartmentState = "MTA"
                    $pool.Open()
                    $runspaces = @()
                    $scriptblock = {
                        Param ($da, $ds, $conn, $queue )
                        $conn.FireInfoMessageEventOnUserErrors = $false
                        $handler = [System.Data.SqlClient.SqlInfoMessageEventHandler] { $queue.Enqueue($_) }
                        $conn.add_InfoMessage($handler)
                        [void]$da.fill($ds)
                    }
                    $queue = New-Object System.Collections.Concurrent.ConcurrentQueue[string]
                    $runspace = [PowerShell]::Create()
                    $null = $runspace.AddScript($scriptblock)
                    $null = $runspace.AddArgument($da)
                    $null = $runspace.AddArgument($ds)
                    $null = $runspace.AddArgument($Conn)
                    $null = $runspace.AddArgument($queue)
                    $runspace.RunspacePool = $pool
                    $runspaces += [PSCustomObject]@{ Pipe = $runspace; Status = $runspace.BeginInvoke() }
                    while ($runspaces.Status.IsCompleted -notcontains $true) {
                        $item = ""
                        if ($queue.TryDequeue([ref]$item)) {
                            "MD this is output streaming: $item"
                        }
                    }
                    foreach ($runspace in $runspaces ) {
                        $results = $runspace.Pipe.EndInvoke($runspace.Status)
                        $runspace.Pipe.Dispose()
                    }
                    $pool.Close()
                    $pool.Dispose()
                }
                else {
                    $conn.FireInfoMessageEventOnUserErrors = $false
                    $handler = [System.Data.SqlClient.SqlInfoMessageEventHandler] { Write-Message -Level Verbose -Message "$_" -FunctionName "QueryMessages" }
                    $conn.add_InfoMessage($handler)
                    [void]$da.fill($ds)
                }
                Write-Message -Level VeryVerbose -Message "end invoke query"
            }
            catch [System.Data.SqlClient.SqlException] {
                # For SQL exception

                $Err = $_

                Write-Message -Level Debug -Message "Capture SQL Error"

                if ($PSBoundParameters.Verbose) {
                    Write-Message -Level Verbose -Message "SQL Error:  $Err"
                } #Shiyang, add the verbose output of exception

                switch ($ErrorActionPreference.tostring()) {
                    { 'SilentlyContinue', 'Ignore' -contains $_ } {

                    }
                    'Stop' {
                        throw $Err
                    }
                    'Continue' {
                        throw $Err
                    }
                    Default {
                        Throw $Err
                    }
                }
            }
            catch {
                # For other exception
                Write-Message -Level Debug -Message "Capture Other Error"

                $Err = $_

                if ($PSBoundParameters.Verbose) {
                    Write-Message -Level Verbose -Message "Other Error:  $Err"
                }

                switch ($ErrorActionPreference.tostring()) {
                    { 'SilentlyContinue', 'Ignore' -contains $_ } {

                    }
                    'Stop' {
                        throw $Err
                    }
                    'Continue' {
                        throw $Err
                    }
                    Default {
                        throw $Err
                    }
                }
            }
            finally {
                #Close the connection
                if (-not $PSBoundParameters.ContainsKey('SQLConnection')) {
                    $Conn.Close()
                }
            }

            if ($AppendServerInstance) {
                #Basics from Chad Miller
                $Column = New-Object Data.DataColumn
                $Column.ColumnName = "ServerInstance"

                if ($ds.Tables.Count -ne 0) {
                    $ds.Tables[0].Columns.Add($Column)
                    Foreach ($row in $ds.Tables[0]) {
                        $row.ServerInstance = $SQLInstance
                    }
                }
            }

            switch ($As) {
                'DataSet' {
                    $ds
                }
                'DataTable' {
                    $ds.Tables
                }
                'DataRow' {
                    if ($ds.Tables.Count -ne 0) {
                        $ds.Tables[0]
                    }
                }
                'PSObject' {
                    if ($ds.Tables.Count -ne 0) {
                        #Scrub DBNulls - Provides convenient results you can use comparisons with
                        #Introduces overhead (e.g. ~2000 rows w/ ~80 columns went from .15 Seconds to .65 Seconds - depending on your data could be much more!)
                        foreach ($row in $ds.Tables[0].Rows) {
                            [DBNullScrubber]::DataRowToPSObject($row)
                        }
                    }
                }
                'SingleValue' {
                    if ($ds.Tables.Count -ne 0) {
                        $ds.Tables[0] | Select-Object -ExpandProperty $ds.Tables[0].Columns[0].ColumnName
                    }
                }
                'Mixed' {
                    $ds
                }
            }
            Write-Message -Level VeryVerbose -Message "end"
        }
    }
}