Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix SqlBulkCopy to work with Data Classification enabled tables #568

Merged
merged 2 commits into from
May 20, 2020
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -2253,40 +2253,40 @@ internal bool TryRun(RunBehavior runBehavior, SqlCommand cmdHandler, SqlDataRead
}
}

if (null != dataStream)
byte peekedToken;
if (!stateObj.TryPeekByte(out peekedToken))
{ // temporarily cache next byte
return false;
}

if (TdsEnums.SQLDATACLASSIFICATION == peekedToken)
{
byte peekedToken;
if (!stateObj.TryPeekByte(out peekedToken))
{ // temporarily cache next byte
byte dataClassificationToken;
if (!stateObj.TryReadByte(out dataClassificationToken))
{
return false;
}
Debug.Assert(TdsEnums.SQLDATACLASSIFICATION == dataClassificationToken);

if (TdsEnums.SQLDATACLASSIFICATION == peekedToken)
SensitivityClassification sensitivityClassification;
if (!TryProcessDataClassification(stateObj, out sensitivityClassification))
{
byte dataClassificationToken;
if (!stateObj.TryReadByte(out dataClassificationToken))
{
return false;
}
Debug.Assert(TdsEnums.SQLDATACLASSIFICATION == dataClassificationToken);

SensitivityClassification sensitivityClassification;
if (!TryProcessDataClassification(stateObj, out sensitivityClassification))
{
return false;
}
if (!dataStream.TrySetSensitivityClassification(sensitivityClassification))
{
return false;
}
return false;
}
if (null != dataStream && !dataStream.TrySetSensitivityClassification(sensitivityClassification))
{
return false;
}

// update peekedToken
if (!stateObj.TryPeekByte(out peekedToken))
{
return false;
}
// update peekedToken
if (!stateObj.TryPeekByte(out peekedToken))
{
return false;
}
}

if (null != dataStream)
{
if (!dataStream.TrySetMetaData(stateObj._cleanupMetaData, (TdsEnums.SQLTABNAME == peekedToken || TdsEnums.SQLCOLINFO == peekedToken)))
{
return false;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -2584,40 +2584,40 @@ internal bool TryRun(RunBehavior runBehavior, SqlCommand cmdHandler, SqlDataRead
}
}

if (null != dataStream)
byte peekedToken;
if (!stateObj.TryPeekByte(out peekedToken))
{ // temporarily cache next byte
return false;
}

if (TdsEnums.SQLDATACLASSIFICATION == peekedToken)
{
byte peekedToken;
if (!stateObj.TryPeekByte(out peekedToken))
{ // temporarily cache next byte
byte dataClassificationToken;
if (!stateObj.TryReadByte(out dataClassificationToken))
{
return false;
}
Debug.Assert(TdsEnums.SQLDATACLASSIFICATION == dataClassificationToken);

if (TdsEnums.SQLDATACLASSIFICATION == peekedToken)
SensitivityClassification sensitivityClassification;
if (!TryProcessDataClassification(stateObj, out sensitivityClassification))
{
byte dataClassificationToken;
if (!stateObj.TryReadByte(out dataClassificationToken))
{
return false;
}
Debug.Assert(TdsEnums.SQLDATACLASSIFICATION == dataClassificationToken);

SensitivityClassification sensitivityClassification;
if (!TryProcessDataClassification(stateObj, out sensitivityClassification))
{
return false;
}
if (!dataStream.TrySetSensitivityClassification(sensitivityClassification))
{
return false;
}
return false;
}
if (null != dataStream && !dataStream.TrySetSensitivityClassification(sensitivityClassification))
{
return false;
}

// update peekedToken
if (!stateObj.TryPeekByte(out peekedToken))
{
return false;
}
// update peekedToken
if (!stateObj.TryPeekByte(out peekedToken))
{
return false;
}
}

if (null != dataStream)
{
if (!dataStream.TrySetMetaData(stateObj._cleanupMetaData, (TdsEnums.SQLTABNAME == peekedToken || TdsEnums.SQLCOLINFO == peekedToken)))
{
return false;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -233,6 +233,31 @@ public static bool IsDatabasePresent(string name)
return present;
}

/// <summary>
/// Checks if object SYS.SENSITIVITY_CLASSIFICATIONS exists in SQL Server
/// </summary>
/// <returns>True, if target SQL Server supports Data Classification</returns>
public static bool IsSupportedDataClassification()
{
try
{
using (var connection = new SqlConnection(TCPConnectionString))
using (var command = new SqlCommand("SELECT * FROM SYS.SENSITIVITY_CLASSIFICATIONS", connection))
{
connection.Open();
command.ExecuteNonQuery();
}
}
catch (SqlException e)
{
// Check for Error 208: Invalid Object Name
if (e.Errors != null && e.Errors[0].Number == 208)
{
return false;
}
}
return true;
}
public static bool IsUdtTestDatabasePresent() => IsDatabasePresent(UdtTestDbName);

public static bool AreConnStringsSetup()
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,7 @@
<Compile Include="AlwaysEncrypted\TestTrustedMasterKeyPaths.cs" />
<Compile Include="DataCommon\AADUtility.cs" />
<Compile Include="DataCommon\CheckConnStrSetupFactAttribute.cs" />
<Compile Include="SQL\DataClassificationTest\DataClassificationTest.cs" />
<Compile Include="TracingTests\EventSourceTest.cs" />
<Compile Include="SQL\AdapterTest\AdapterTest.cs" />
<Compile Include="SQL\AsyncTest\BeginExecAsyncTest.cs" />
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,209 @@
// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.
// See the LICENSE file in the project root for more information.

using System;
using System.Collections.ObjectModel;
using System.Data;
using Microsoft.Data.SqlClient.DataClassification;
using Xunit;

namespace Microsoft.Data.SqlClient.ManualTesting.Tests
{
public static class DataClassificationTest
{
private static string s_tableName;

[ConditionalFact(typeof(DataTestUtility), nameof(DataTestUtility.AreConnStringsSetup), nameof(DataTestUtility.IsNotAzureServer), nameof(DataTestUtility.IsSupportedDataClassification))]
public static void TestDataClassificationResultSet()
{
s_tableName = DataTestUtility.GetUniqueNameForSqlServer("DC");
using (SqlConnection sqlConnection = new SqlConnection(DataTestUtility.TCPConnectionString))
using (SqlCommand sqlCommand = sqlConnection.CreateCommand())
{
try
{
sqlConnection.Open();
Assert.True(DataTestUtility.IsSupportedDataClassification());
CreateTable(sqlCommand);
RunTestsForServer(sqlCommand);
}
finally
{
DataTestUtility.DropTable(sqlConnection, s_tableName);
}
}
}

private static void RunTestsForServer(SqlCommand sqlCommand)
{
sqlCommand.CommandText = "SELECT * FROM " + s_tableName;
using (SqlDataReader reader = sqlCommand.ExecuteReader())
{
VerifySensitivityClassification(reader);
}
}

private static void VerifySensitivityClassification(SqlDataReader reader)
{
if (null != reader.SensitivityClassification)
{
for (int columnPos = 0; columnPos < reader.SensitivityClassification.ColumnSensitivities.Count;
columnPos++)
{
foreach (SensitivityProperty sp in reader.SensitivityClassification.ColumnSensitivities[columnPos].SensitivityProperties)
{
ReadOnlyCollection<InformationType> infoTypes = reader.SensitivityClassification.InformationTypes;
Assert.Equal(3, infoTypes.Count);
for (int i = 0; i < infoTypes.Count; i++)
{
VerifyInfoType(infoTypes[i], i + 1);
}

ReadOnlyCollection<Label> labels = reader.SensitivityClassification.Labels;
Assert.Single(labels);
VerifyLabel(labels[0]);

if (columnPos == 1 || columnPos == 2 || columnPos == 6 || columnPos == 7)
{
VerifyLabel(sp.Label);
VerifyInfoType(sp.InformationType, columnPos);
}
}
}
}
}

private static void VerifyLabel(Label label)
{
Assert.True(label != null);
Assert.Equal("L1", label.Id);
Assert.Equal("PII", label.Name);
}

private static void VerifyInfoType(InformationType informationType, int i)
{
Assert.True(informationType != null);
Assert.Equal(i == 1 ? "COMPANY" : (i == 2 ? "NAME" : "CONTACT"), informationType.Id);
Assert.Equal(i == 1 ? "Company Name" : (i == 2 ? "Person Name" : "Contact Information"), informationType.Name);
}

private static void CreateTable(SqlCommand sqlCommand)
{
sqlCommand.CommandText = "CREATE TABLE " + s_tableName + " ("
+ "[Id] [int] IDENTITY(1,1) NOT NULL,"
+ "[CompanyName] [nvarchar](40) NOT NULL,"
+ "[ContactName] [nvarchar](50) NULL,"
+ "[ContactTitle] [nvarchar](40) NULL,"
+ "[City] [nvarchar](40) NULL,"
+ "[CountryName] [nvarchar](40) NULL,"
+ "[Phone] [nvarchar](30) MASKED WITH (FUNCTION = 'default()') NULL,"
+ "[Fax] [nvarchar](30) MASKED WITH (FUNCTION = 'default()') NULL)";
sqlCommand.ExecuteNonQuery();

sqlCommand.CommandText = "ADD SENSITIVITY CLASSIFICATION TO " + s_tableName
+ ".CompanyName WITH (LABEL='PII', LABEL_ID='L1', INFORMATION_TYPE='Company Name', INFORMATION_TYPE_ID='COMPANY')";
sqlCommand.ExecuteNonQuery();

sqlCommand.CommandText = "ADD SENSITIVITY CLASSIFICATION TO " + s_tableName
+ ".ContactName WITH (LABEL='PII', LABEL_ID='L1', INFORMATION_TYPE='Person Name', INFORMATION_TYPE_ID='NAME')";
sqlCommand.ExecuteNonQuery();

sqlCommand.CommandText = "ADD SENSITIVITY CLASSIFICATION TO " + s_tableName
+ ".Phone WITH (LABEL='PII', LABEL_ID='L1', INFORMATION_TYPE='Contact Information', INFORMATION_TYPE_ID='CONTACT')";
sqlCommand.ExecuteNonQuery();

sqlCommand.CommandText = "ADD SENSITIVITY CLASSIFICATION TO " + s_tableName
+ ".Fax WITH (LABEL='PII', LABEL_ID='L1', INFORMATION_TYPE='Contact Information', INFORMATION_TYPE_ID='CONTACT')";
sqlCommand.ExecuteNonQuery();

// INSERT ROWS OF DATA
sqlCommand.CommandText = "INSERT INTO " + s_tableName + " VALUES (@companyName, @contactName, @contactTitle, @city, @country, @phone, @fax)";

sqlCommand.Parameters.AddWithValue("@companyName", "Exotic Liquids");
sqlCommand.Parameters.AddWithValue("@contactName", "Charlotte Cooper");
sqlCommand.Parameters.AddWithValue("@contactTitle", "");
sqlCommand.Parameters.AddWithValue("city", "London");
sqlCommand.Parameters.AddWithValue("@country", "UK");
sqlCommand.Parameters.AddWithValue("@phone", "(171) 555-2222");
sqlCommand.Parameters.AddWithValue("@fax", "");
sqlCommand.ExecuteNonQuery();

sqlCommand.Parameters.Clear();
sqlCommand.Parameters.AddWithValue("@companyName", "New Orleans");
sqlCommand.Parameters.AddWithValue("@contactName", "Cajun Delights");
sqlCommand.Parameters.AddWithValue("@contactTitle", "");
sqlCommand.Parameters.AddWithValue("city", "New Orleans");
sqlCommand.Parameters.AddWithValue("@country", "USA");
sqlCommand.Parameters.AddWithValue("@phone", "(100) 555-4822");
sqlCommand.Parameters.AddWithValue("@fax", "");
sqlCommand.ExecuteNonQuery();

sqlCommand.Parameters.Clear();
sqlCommand.Parameters.AddWithValue("@companyName", "Grandma Kelly's Homestead");
sqlCommand.Parameters.AddWithValue("@contactName", "Regina Murphy");
sqlCommand.Parameters.AddWithValue("@contactTitle", "");
sqlCommand.Parameters.AddWithValue("@city", "Ann Arbor");
sqlCommand.Parameters.AddWithValue("@country", "USA");
sqlCommand.Parameters.AddWithValue("@phone", "(313) 555-5735");
sqlCommand.Parameters.AddWithValue("@fax", "(313) 555-3349");
sqlCommand.ExecuteNonQuery();
}

[ConditionalFact(typeof(DataTestUtility), nameof(DataTestUtility.AreConnStringsSetup), nameof(DataTestUtility.IsNotAzureServer), nameof(DataTestUtility.IsSupportedDataClassification))]
public static void TestDataClassificationBulkCopy()
{
var data = new DataTable("Company");
data.Columns.Add("CompanyId", typeof(Guid));
data.Columns.Add("CompanyName", typeof(string));
data.Columns.Add("Email", typeof(string));
data.Columns.Add("CompanyType", typeof(int));

data.Rows.Add(Guid.NewGuid(), "Company 1", "sample1@contoso.com", 1);
data.Rows.Add(Guid.NewGuid(), "Company 2", "sample2@contoso.com", 1);
data.Rows.Add(Guid.NewGuid(), "Company 3", "sample3@contoso.com", 1);

var tableName = DataTestUtility.GetUniqueNameForSqlServer("DC");

using (var connection = new SqlConnection(DataTestUtility.TCPConnectionString))
{
connection.Open();
try
{
// Setup Table
using (SqlCommand sqlCommand = connection.CreateCommand())
{
sqlCommand.CommandText = $"CREATE TABLE {tableName} (" +
$" [CompanyId] [uniqueidentifier] NOT NULL," +
$" [CompanyName][nvarchar](255) NOT NULL," +
$" [Email] [nvarchar](50) NULL," +
$" [CompanyType] [int] not null," +
$" CONSTRAINT[PK_Company] PRIMARY KEY CLUSTERED (" +
$" [CompanyId] ASC" +
$" ) WITH(STATISTICS_NORECOMPUTE = OFF, IGNORE_DUP_KEY = OFF) ON[PRIMARY]" +
$" ) ON[PRIMARY]";
sqlCommand.ExecuteNonQuery();

sqlCommand.CommandText = $"ADD SENSITIVITY CLASSIFICATION TO {tableName}.Email WITH (label = 'Confidential', label_id = 'c185460f-4e20-4b89-9876-ae95f07ba087', information_type = 'Contact Info', information_type_id = '5c503e21-22c6-81fa-620b-f369b8ec38d1');";
sqlCommand.ExecuteNonQuery();
}

// Perform Bulk Insert
using (var bulk = new SqlBulkCopy(connection))
{
bulk.DestinationTableName = tableName;
bulk.ColumnMappings.Add("CompanyId", "CompanyId");
bulk.ColumnMappings.Add("CompanyName", "CompanyName");
bulk.ColumnMappings.Add("Email", "Email");
bulk.ColumnMappings.Add("CompanyType", "CompanyType");
bulk.WriteToServer(data);
}
}
finally
{
DataTestUtility.DropTable(connection, tableName);
}
}
}
}
}