Skip to content

Commit e79bef9

Browse files
andreasjordanandreasjordan
authored
Invoke-DbaDbDataMasking: Fix MaskingID column and index left behind after masking (#10223)
1 parent 23b429b commit e79bef9

1 file changed

Lines changed: 81 additions & 68 deletions

File tree

‎public/Invoke-DbaDbDataMasking.ps1‎

Lines changed: 81 additions & 68 deletions
Original file line numberDiff line numberDiff line change
@@ -352,58 +352,68 @@ function Invoke-DbaDbDataMasking {
352352

353353
[bool]$cleanupIdentityColumn = $false
354354

355+
# The masking index name used for cleanup checks
356+
$maskingIndexName = "NIX__$($dbTable.Schema)_$($dbTable.Name)_Masking"
357+
355358
# Make sure there is an identity column present to speed things up
356-
if (-not ($dbTable.Columns | Where-Object { $_.Identity -eq $true })) {
357-
Write-Message -Level Verbose -Message "Adding identity column to table [$($dbTable.Schema)].[$($dbTable.Name)]"
358-
$query = "ALTER TABLE [$($dbTable.Schema)].[$($dbTable.Name)] ADD MaskingID BIGINT IDENTITY(1, 1) NOT NULL;"
359+
# Skip column and index creation when -WhatIf is active to avoid leaving behind schema changes
360+
if (-not $WhatIfPreference) {
361+
if (-not ($dbTable.Columns | Where-Object { $_.Identity -eq $true })) {
362+
Write-Message -Level Verbose -Message "Adding identity column to table [$($dbTable.Schema)].[$($dbTable.Name)]"
363+
$query = "ALTER TABLE [$($dbTable.Schema)].[$($dbTable.Name)] ADD MaskingID BIGINT IDENTITY(1, 1) NOT NULL;"
359364

360-
try {
361-
Invoke-DbaQuery -SqlInstance $server -SqlCredential $SqlCredential -Database $db.Name -Query $query
362-
} catch {
363-
Stop-Function -Message "Could not alter the table to add the masking id" -Target $db -Continue
364-
}
365+
try {
366+
Invoke-DbaQuery -SqlInstance $server -SqlCredential $SqlCredential -Database $db.Name -Query $query
367+
} catch {
368+
Stop-Function -Message "Could not alter the table to add the masking id" -Target $db -Continue
369+
}
365370

366-
$cleanupIdentityColumn = $true
371+
$cleanupIdentityColumn = $true
367372

368-
$identityColumn = "MaskingID"
373+
$identityColumn = "MaskingID"
369374

370-
$dbTable.Columns.Refresh()
371-
} else {
372-
$identityColumn = $dbTable.Columns | Where-Object { $_.Identity } | Select-Object -ExpandProperty Name
373-
}
375+
$dbTable.Columns.Refresh()
376+
} else {
377+
$identityColumn = $dbTable.Columns | Where-Object { $_.Identity } | Select-Object -ExpandProperty Name
378+
}
374379

375-
# Check if the index for the identity column is already present
376-
$maskingIndexName = "NIX__$($dbTable.Schema)_$($dbTable.Name)_Masking"
377-
try {
378-
if ($dbTable.Indexes.Name -contains $maskingIndexName) {
379-
Write-Message -Level Verbose -Message "Masking index already exists in table [$($dbTable.Schema)].[$($dbTable.Name)]. Dropping it..."
380-
$dbTable.Indexes[$($maskingIndexName)].Drop()
380+
# Check if the index for the identity column is already present
381+
try {
382+
if ($dbTable.Indexes.Name -contains $maskingIndexName) {
383+
Write-Message -Level Verbose -Message "Masking index already exists in table [$($dbTable.Schema)].[$($dbTable.Name)]. Dropping it..."
384+
$dbTable.Indexes[$($maskingIndexName)].Drop()
385+
}
386+
} catch {
387+
Stop-Function -Message "Could not remove identity index to table [$($dbTable.Schema)].[$($dbTable.Name)]" -Continue
381388
}
382-
} catch {
383-
Stop-Function -Message "Could not remove identity index to table [$($dbTable.Schema)].[$($dbTable.Name)]" -Continue
384-
}
385389

386-
# Create the index for the identity column
387-
try {
388-
Write-Message -Level Verbose -Message "Adding index on identity column [$($identityColumn)] in table [$($dbTable.Schema)].[$($dbTable.Name)]"
390+
# Create the index for the identity column
391+
try {
392+
Write-Message -Level Verbose -Message "Adding index on identity column [$($identityColumn)] in table [$($dbTable.Schema)].[$($dbTable.Name)]"
389393

390-
$query = "CREATE NONCLUSTERED INDEX [$($maskingIndexName)] ON [$($dbTable.Schema)].[$($dbTable.Name)]([$($identityColumn)])"
394+
$query = "CREATE NONCLUSTERED INDEX [$($maskingIndexName)] ON [$($dbTable.Schema)].[$($dbTable.Name)]([$($identityColumn)])"
391395

392-
$queryParams = @{
393-
SqlInstance = $server
394-
SqlCredential = $SqlCredential
395-
Database = $db.Name
396-
Query = $query
397-
QueryTimeout = $CommandTimeout
398-
}
396+
$queryParams = @{
397+
SqlInstance = $server
398+
SqlCredential = $SqlCredential
399+
Database = $db.Name
400+
Query = $query
401+
QueryTimeout = $CommandTimeout
402+
}
399403

400-
Invoke-DbaQuery @queryParams
401-
} catch {
402-
Stop-Function -Message "Could not add identity index to table [$($dbTable.Schema)].[$($dbTable.Name)]" -Continue
404+
Invoke-DbaQuery @queryParams
405+
} catch {
406+
Stop-Function -Message "Could not add identity index to table [$($dbTable.Schema)].[$($dbTable.Name)]" -Continue
407+
}
403408
}
404409

405410
try {
406-
if (-not $tableobject.FilterQuery) {
411+
if ($WhatIfPreference) {
412+
# In WhatIf mode, only get the row count without modifying the table structure
413+
$query = "SELECT COUNT(*) AS RowCount FROM [$($tableobject.Schema)].[$($tableobject.Name)]"
414+
$rowCount = ($db.Query($query)).RowCount
415+
$data = New-Object object[] $rowCount
416+
} elseif (-not $tableobject.FilterQuery) {
407417
# Get all the columns from the table
408418
$columnString = "[" + (($dbTable.Columns | Where-Object { $_.DataType -in $supportedDataTypes } | Select-Object Name -ExpandProperty Name) -join "],[") + "]"
409419

@@ -412,6 +422,9 @@ function Invoke-DbaDbDataMasking {
412422

413423
# Put it all together
414424
$query = "SELECT $($columnString) FROM [$($tableobject.Schema)].[$($tableobject.Name)]"
425+
426+
# Get the data
427+
[array]$data = $db.Query($query)
415428
} else {
416429
# Get the query from the table objects
417430
$query = ($tableobject.FilterQuery).ToLower()
@@ -424,10 +437,10 @@ function Invoke-DbaDbDataMasking {
424437
# Put it all together again with the identifier
425438
$query = "$($queryParts[0].Trim()), $($identityColumn) FROM $($queryParts[1].Trim())"
426439
}
427-
}
428440

429-
# Get the data
430-
[array]$data = $db.Query($query)
441+
# Get the data
442+
[array]$data = $db.Query($query)
443+
}
431444
} catch {
432445
Stop-Function -Message "Failure retrieving the data from table [$($tableobject.Schema)].[$($tableobject.Name)]" -Target $Database -ErrorRecord $_ -Continue
433446
}
@@ -1232,33 +1245,6 @@ function Invoke-DbaDbDataMasking {
12321245
$null = $stringBuilder.Clear()
12331246
}
12341247

1235-
# Clean up the masking index
1236-
try {
1237-
# Refresh the indexes to make sure to have the latest list
1238-
$dbTable.Indexes.Refresh()
1239-
1240-
# Check if the index is there
1241-
if ($dbTable.Indexes.Name -contains $maskingIndexName) {
1242-
Write-Message -Level verbose -Message "Removing identity index from table [$($dbTable.Schema)].[$($dbTable.Name)]"
1243-
$dbTable.Indexes[$($maskingIndexName)].Drop()
1244-
}
1245-
} catch {
1246-
Stop-Function -Message "Could not remove identity index from table [$($dbTable.Schema)].[$($dbTable.Name)]" -Continue
1247-
}
1248-
1249-
# Clean up the identity column
1250-
if ($cleanupIdentityColumn) {
1251-
try {
1252-
Write-Message -Level Verbose -Message "Removing identity column [$($identityColumn)] from table [$($dbTable.Schema)].[$($dbTable.Name)]"
1253-
1254-
$query = "ALTER TABLE [$($dbTable.Schema)].[$($dbTable.Name)] DROP COLUMN [$($identityColumn)]"
1255-
1256-
Invoke-DbaQuery -SqlInstance $instance -SqlCredential $SqlCredential -Database $db.Name -Query $query -EnableException
1257-
} catch {
1258-
Stop-Function -Message "Could not remove identity column from table [$($dbTable.Schema)].[$($dbTable.Name)]" -Continue
1259-
}
1260-
}
1261-
12621248
# Return the masking results
12631249
if ($maskingErrorFlag) {
12641250
$maskingStatus = "Failed"
@@ -1284,6 +1270,33 @@ function Invoke-DbaDbDataMasking {
12841270
$null = $elapsed.Reset()
12851271
}
12861272

1273+
# Clean up the masking index (always runs, regardless of -WhatIf or errors during masking)
1274+
try {
1275+
# Refresh the indexes to make sure to have the latest list
1276+
$dbTable.Indexes.Refresh()
1277+
1278+
# Check if the index is there
1279+
if ($dbTable.Indexes.Name -contains $maskingIndexName) {
1280+
Write-Message -Level verbose -Message "Removing identity index from table [$($dbTable.Schema)].[$($dbTable.Name)]"
1281+
$dbTable.Indexes[$($maskingIndexName)].Drop()
1282+
}
1283+
} catch {
1284+
Stop-Function -Message "Could not remove identity index from table [$($dbTable.Schema)].[$($dbTable.Name)]" -Continue
1285+
}
1286+
1287+
# Clean up the identity column (always runs, regardless of -WhatIf or errors during masking)
1288+
if ($cleanupIdentityColumn) {
1289+
try {
1290+
Write-Message -Level Verbose -Message "Removing identity column [$($identityColumn)] from table [$($dbTable.Schema)].[$($dbTable.Name)]"
1291+
1292+
$query = "ALTER TABLE [$($dbTable.Schema)].[$($dbTable.Name)] DROP COLUMN [$($identityColumn)]"
1293+
1294+
Invoke-DbaQuery -SqlInstance $instance -SqlCredential $SqlCredential -Database $db.Name -Query $query -EnableException
1295+
} catch {
1296+
Stop-Function -Message "Could not remove identity column from table [$($dbTable.Schema)].[$($dbTable.Name)]" -Continue
1297+
}
1298+
}
1299+
12871300
# Cleanup
12881301
if ($uniqueDataTableName) {
12891302
Write-Message -Message "Cleaning up unique temporary table '$uniqueDataTableName'" -Level verbose

0 commit comments

Comments
 (0)