Automate Range Right Partition Management with PowerShell

###############################################################################
# File:        SwitchPartition.ps1
# Author:     Hugh Scott
# Date:        2009/12/03
# 
# Description:
# Script switches data from a partitioned table to a staging table in preparation
# for dropping the staging table. Script identifies all tables associated to 
# a particular PartitionScheme, builds the necessary staging table, switches the
# partition and then (optionally) drops the staging table(s).
#
# ToDo:
# 1. Add a switch to allow the user to optionally specify whether or not to perform
# a merge (remove that partition from the partition scheme)
# 2. Convert $partitionNumber to $boundaryValue
#
# Usage:
#
# .\SwitchPartition.ps1 Server Database stageTableSuffix partitionSchemeName boundaryDate [-keepStageTable] [-keepPartition] [-keepFileGroup]
#
###############################################################################
param (
    [string]$serverName ,         
    [string]$databaseName ,     
    [string]$stageTableSuffix ,     
    [string]$partitionSchemeName ,
    [datetime]$boundaryDate ,
    [switch]$keepStageTable ,
    [switch]$keepPartition ,
    [switch]$keepFileGroup
)

## Logical check of Parameters:
## 1. If $keepFileGroup is empty, then $keepStageTable and $keepPartition must also be empty
## 2. If $keepPartition is empty, then $keepStageTable must also be empty
## 3. Validate $server
## 4. Validate $databaseName
## 5. Validate $partitionSchemeName
if (!$keepFileGroup -and ($keepStageTable -or $keepPartition))
{
    Throw("Error: If file group is to be dropped, stage table(s) must be dropped and partition merged.")
}
if (!$keepPartition -and $keepStageTable)
{
    Throw("Error: If partition is to be merged, stage table(s) must be dropped.")
}

## Create a scripting object, specify options and then execute:
[System.Reflection.Assembly]::LoadWithPartialName("Microsoft.SqlServer.SMO") | Out-Null
$srv = New-Object ("Microsoft.SqlServer.Management.SMO.Server")($serverName)
$db = $srv.Databases[$databaseName]
if($db.Name -eq $Null)
{
    Throw("Error: Unable to connect to database $databaseName on server $serverName.")
}

## Validate that $partitionSchemeName exists
$ps = $db.PartitionSchemes[$partitionSchemeName]
if($ps.Name -eq $Null)
{
    Throw("Error: Partition Scheme $partitionSchemeName does not exist")
}

## Determine partitionNumber from boundaryDate, spin through the rangeValues
## if someone has a better method, please let me know!
$pf = $db.PartitionFunctions[$ps.PartitionFunction]
$rangeValues = $pf.RangeValues
$i = 1
foreach($Value in $rangeValues)
{
    if($Value -eq $boundaryDate)
    {
        $partitionNumber = $i
    }
    $i = $i + 1
}

## Validate the resulting $partitionNumber value
## The resulting $partitionNumber can't be null
if($partitionNumber -eq $Null)
{
    Throw ("Error: Invalid boundary date $boundaryDate!")
}
else
{
    # If this is a range-right scheme, we have to increment the partition number by 1
    # At least, I THINK this is why we have to increment it by one
    if ($pf.RangeType -eq "Right")
    {
        $partitionNumber = $partitionNumber + 1
    }
    
    # $partitionNumber must be > 1 and < max number of partitions
    if($partitionNumber -le 1 -or $partitionNumber -ge $pf.NumberOfPartitions )
    {
    Throw ("Error: Invalid boundary date $boundaryDate!")
    }
}

## Create a collection of all tables associated to the PartitionScheme
## 1. Iterate through the table collection, 
## 2. Create a copy of the table,
## 3. Switch the data from the specified partition into the table 
## 4. Drop the copy of the table
$partitionTables = $db.tables | Where-Object {$_.PartitionScheme -eq $partitionSchemeName}
foreach ($partitionTable in $partitionTables)
{
    $stageSchemaName = $partitionTable.Schema
    $stageTableName = $partitionTable.Name + $stageTableSuffix
    # Write-Host $stageTableName

    # Test for the existence of the staging table, if it exists, drop it and then create it
    $testTable = $db.tables | Where-Object {$_.schema -eq $stageSchemaName-and$_.name -eq $stageTableName}
    If($testTable -ne $Null)
    {
        $testTable.Drop()
    }
    # Create the staging table
    $stageTable = New-Object("Microsoft.SqlServer.Management.SMO.Table")($db, $stageTableName, $stageSchemaName)

    # Create a copy of the partitioned table
    foreach    ($Col in $partitionTable.Columns)
    {
        $stageCol = New-Object("Microsoft.SqlServer.Management.SMO.Column")($stageTable, $Col.name, $Col.datatype)
        $stageCol.Collation = $Col.Collation
        $stageCol.Nullable = $Col.Nullable
        $stageCol.ComputedText = $Col.ComputedText
        $stageCol.Computed = $Col.Computed
        $stageCol.Default = $Col.Default
    #    if($Col.DefaultConstraint -ne $null)
    #    {
    #        $stageCol.AddDefaultConstraint($stageTable.name + "_" + $Col.DefaultConstraint.Name)
    #        $stageCol.DefaultConstraint.Text = $Col.DefaultConstraint.Text        
    #    }
        $stageCol.IsPersisted = $Col.IsPersisted
        $stageCol.DefaultSchema = $Col.DefaultSchema
        $stageCol.RowGuidCol = $Col.RowGuidCol
        
        $stageTable.Columns.Add($stageCol)
    }
    
    $stageTable.AnsiNullsStatus = $partitionTable.AnsiNullsStatus
    $stageTable.QuotedIdentifierStatus = $partitionTable.QuotedIdentifierStatus
    $stageTable.TextFileGroup = $partitionTable.TextFileGroup
    $stageTable.FileGroup = $db.PartitionSchemes[$partitionTable.PartitionScheme].FileGroups[$partitionNumber - 1]
    
    $stageTable.Create()

    # If it has one, create the clustered index on the table
    if($partitionTable.HasClusteredIndex)
    {
        foreach ($index in $partitionTable.Indexes)
        {
         if ($index.IsClustered)
         {
             $stageIndexName = $stageTable.Name + "_" + $index.Name
                $stgIndex = New-Object("Microsoft.SqlServer.Management.SMO.Index")($stageTable, $stageIndexName)
                foreach ($iCol in $index.IndexedColumns)
                {
                 $stgICol = New-Object("Microsoft.SqlServer.Management.SMO.IndexedColumn")($stgIndex, $iCol.Name, $iCol.Descending)
                 $stgICol.IsIncluded = $iCol.IsIncluded
                 $stgIndex.IndexedColumns.Add($stgICol)
                }
                $stgIndex.IndexKeyType = $index.IndexKeyType
                $stgIndex.IsClustered = $index.IsClustered
                $stgIndex.IsUnique = $index.IsUnique
                $stgIndex.CompactLargeObjects = $index.CompactLargeObjects
                $stgIndex.IgnoreDuplicateKeys = $index.IgnoreDuplicateKeys
                $stgIndex.IsFullTextKey = $index.IsFullTextKey
                $stgIndex.PadIndex = $index.PadIndex
                $stgIndex.FileGroup = $db.PartitionSchemes[$index.PartitionScheme].FileGroups[$partitionNumber - 1]
                
                ## add the partitioning column to the index if it is not already there
                $partitionKeyName = $index.PartitionSchemeParameters[0].Name
                if ($stgIndex.IndexedColumns[$partitionKeyName] -eq $Null)
                {
                 $stgICol = New-Object("Microsoft.SqlServer.Management.SMO.IndexedColumn")($stgIndex, $partitionKeyName)
                 ## It is added as a Key to the Clustered index and as an Include column to a Nonclustered
                 $stgICol.IsIncluded = !$stgIndex.IsClustered
                 $stgIndex.IndexedColumns.Add($stgICol)
                }
                
                $stgIndex.Create()
         # break
         }
        }        
    }
    
    # Switch partition (put the data in the indicated partition into the staging table
    $partitionTable.SwitchPartition($partitionNumber, $stageTable)
    
    # If keepTable is specified, then keep the staging table (for testing/validation
    # Otherwise, drop it
    if(!$keepStageTable)
    {
        $stageTable.Drop()
    }    
    else
    {
        Write-Host "Task not performed: Drop table $stageTable.ToString()"
    }
}

## Alter partition function; merge the specified boundary
if(!$keepPartition)
{
    $pf.MergeRangePartition($boundaryDate)
}
else
{
    Write-Host "Task not performed: Merge Range Partition $boundaryDate"
}

## If specified, drop filegroup from the database 
if(!$keepFileGroup)
{
    $fileGroupName = $db.PartitionSchemes[$partitionSchemeName].FileGroups[$partitionNumber - 1]
    $fileGroup = $db.FileGroups[$fileGroupName]
    $fileGroup.Drop()
}
else
{
    $fileGroup = $db.PartitionSchemes[$partitionSchemeName].FileGroups[$partitionNumber - 1]
    Write-Host "Task not performed: Drop filegroup name $fileGroup.Name"
}