Skip to content

Commit 1395901

Browse files
cheenamalhotracheenamalhotra
authored
Fix SqlBulkCopy to work with Data Classification enabled tables (#568)
1 parent 7ebb96d commit 1395901

File tree

5 files changed

+287
-52
lines changed

5 files changed

+287
-52
lines changed

‎src/Microsoft.Data.SqlClient/netcore/src/Microsoft/Data/SqlClient/TdsParser.cs‎

Lines changed: 26 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -2253,40 +2253,40 @@ internal bool TryRun(RunBehavior runBehavior, SqlCommand cmdHandler, SqlDataRead
22532253
}
22542254
}
22552255

2256-
if (null != dataStream)
2256+
byte peekedToken;
2257+
if (!stateObj.TryPeekByte(out peekedToken))
2258+
{ // temporarily cache next byte
2259+
return false;
2260+
}
2261+
2262+
if (TdsEnums.SQLDATACLASSIFICATION == peekedToken)
22572263
{
2258-
byte peekedToken;
2259-
if (!stateObj.TryPeekByte(out peekedToken))
2260-
{ // temporarily cache next byte
2264+
byte dataClassificationToken;
2265+
if (!stateObj.TryReadByte(out dataClassificationToken))
2266+
{
22612267
return false;
22622268
}
2269+
Debug.Assert(TdsEnums.SQLDATACLASSIFICATION == dataClassificationToken);
22632270

2264-
if (TdsEnums.SQLDATACLASSIFICATION == peekedToken)
2271+
SensitivityClassification sensitivityClassification;
2272+
if (!TryProcessDataClassification(stateObj, out sensitivityClassification))
22652273
{
2266-
byte dataClassificationToken;
2267-
if (!stateObj.TryReadByte(out dataClassificationToken))
2268-
{
2269-
return false;
2270-
}
2271-
Debug.Assert(TdsEnums.SQLDATACLASSIFICATION == dataClassificationToken);
2272-
2273-
SensitivityClassification sensitivityClassification;
2274-
if (!TryProcessDataClassification(stateObj, out sensitivityClassification))
2275-
{
2276-
return false;
2277-
}
2278-
if (!dataStream.TrySetSensitivityClassification(sensitivityClassification))
2279-
{
2280-
return false;
2281-
}
2274+
return false;
2275+
}
2276+
if (null != dataStream && !dataStream.TrySetSensitivityClassification(sensitivityClassification))
2277+
{
2278+
return false;
2279+
}
22822280

2283-
// update peekedToken
2284-
if (!stateObj.TryPeekByte(out peekedToken))
2285-
{
2286-
return false;
2287-
}
2281+
// update peekedToken
2282+
if (!stateObj.TryPeekByte(out peekedToken))
2283+
{
2284+
return false;
22882285
}
2286+
}
22892287

2288+
if (null != dataStream)
2289+
{
22902290
if (!dataStream.TrySetMetaData(stateObj._cleanupMetaData, (TdsEnums.SQLTABNAME == peekedToken || TdsEnums.SQLCOLINFO == peekedToken)))
22912291
{
22922292
return false;

‎src/Microsoft.Data.SqlClient/netfx/src/Microsoft/Data/SqlClient/TdsParser.cs‎

Lines changed: 26 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -2584,40 +2584,40 @@ internal bool TryRun(RunBehavior runBehavior, SqlCommand cmdHandler, SqlDataRead
25842584
}
25852585
}
25862586

2587-
if (null != dataStream)
2587+
byte peekedToken;
2588+
if (!stateObj.TryPeekByte(out peekedToken))
2589+
{ // temporarily cache next byte
2590+
return false;
2591+
}
2592+
2593+
if (TdsEnums.SQLDATACLASSIFICATION == peekedToken)
25882594
{
2589-
byte peekedToken;
2590-
if (!stateObj.TryPeekByte(out peekedToken))
2591-
{ // temporarily cache next byte
2595+
byte dataClassificationToken;
2596+
if (!stateObj.TryReadByte(out dataClassificationToken))
2597+
{
25922598
return false;
25932599
}
2600+
Debug.Assert(TdsEnums.SQLDATACLASSIFICATION == dataClassificationToken);
25942601

2595-
if (TdsEnums.SQLDATACLASSIFICATION == peekedToken)
2602+
SensitivityClassification sensitivityClassification;
2603+
if (!TryProcessDataClassification(stateObj, out sensitivityClassification))
25962604
{
2597-
byte dataClassificationToken;
2598-
if (!stateObj.TryReadByte(out dataClassificationToken))
2599-
{
2600-
return false;
2601-
}
2602-
Debug.Assert(TdsEnums.SQLDATACLASSIFICATION == dataClassificationToken);
2603-
2604-
SensitivityClassification sensitivityClassification;
2605-
if (!TryProcessDataClassification(stateObj, out sensitivityClassification))
2606-
{
2607-
return false;
2608-
}
2609-
if (!dataStream.TrySetSensitivityClassification(sensitivityClassification))
2610-
{
2611-
return false;
2612-
}
2605+
return false;
2606+
}
2607+
if (null != dataStream && !dataStream.TrySetSensitivityClassification(sensitivityClassification))
2608+
{
2609+
return false;
2610+
}
26132611

2614-
// update peekedToken
2615-
if (!stateObj.TryPeekByte(out peekedToken))
2616-
{
2617-
return false;
2618-
}
2612+
// update peekedToken
2613+
if (!stateObj.TryPeekByte(out peekedToken))
2614+
{
2615+
return false;
26192616
}
2617+
}
26202618

2619+
if (null != dataStream)
2620+
{
26212621
if (!dataStream.TrySetMetaData(stateObj._cleanupMetaData, (TdsEnums.SQLTABNAME == peekedToken || TdsEnums.SQLCOLINFO == peekedToken)))
26222622
{
26232623
return false;

‎src/Microsoft.Data.SqlClient/tests/ManualTests/DataCommon/DataTestUtility.cs‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -233,6 +233,31 @@ public static bool IsDatabasePresent(string name)
233233
return present;
234234
}
235235

236+
/// <summary>
237+
/// Checks if object SYS.SENSITIVITY_CLASSIFICATIONS exists in SQL Server
238+
/// </summary>
239+
/// <returns>True, if target SQL Server supports Data Classification</returns>
240+
public static bool IsSupportedDataClassification()
241+
{
242+
try
243+
{
244+
using (var connection = new SqlConnection(TCPConnectionString))
245+
using (var command = new SqlCommand("SELECT * FROM SYS.SENSITIVITY_CLASSIFICATIONS", connection))
246+
{
247+
connection.Open();
248+
command.ExecuteNonQuery();
249+
}
250+
}
251+
catch (SqlException e)
252+
{
253+
// Check for Error 208: Invalid Object Name
254+
if (e.Errors != null && e.Errors[0].Number == 208)
255+
{
256+
return false;
257+
}
258+
}
259+
return true;
260+
}
236261
public static bool IsUdtTestDatabasePresent() => IsDatabasePresent(UdtTestDbName);
237262

238263
public static bool AreConnStringsSetup()

‎src/Microsoft.Data.SqlClient/tests/ManualTests/Microsoft.Data.SqlClient.ManualTesting.Tests.csproj‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,7 @@
5959
<Compile Include="AlwaysEncrypted\TestTrustedMasterKeyPaths.cs" />
6060
<Compile Include="DataCommon\AADUtility.cs" />
6161
<Compile Include="DataCommon\CheckConnStrSetupFactAttribute.cs" />
62+
<Compile Include="SQL\DataClassificationTest\DataClassificationTest.cs" />
6263
<Compile Include="TracingTests\EventSourceTest.cs" />
6364
<Compile Include="SQL\AdapterTest\AdapterTest.cs" />
6465
<Compile Include="SQL\AsyncTest\BeginExecAsyncTest.cs" />

‎src/Microsoft.Data.SqlClient/tests/ManualTests/SQL/DataClassificationTest/DataClassificationTest.cs‎

Lines changed: 209 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,209 @@
1+
// Licensed to the .NET Foundation under one or more agreements.
2+
// The .NET Foundation licenses this file to you under the MIT license.
3+
// See the LICENSE file in the project root for more information.
4+
5+
using System;
6+
using System.Collections.ObjectModel;
7+
using System.Data;
8+
using Microsoft.Data.SqlClient.DataClassification;
9+
using Xunit;
10+
11+
namespace Microsoft.Data.SqlClient.ManualTesting.Tests
12+
{
13+
public static class DataClassificationTest
14+
{
15+
private static string s_tableName;
16+
17+
[ConditionalFact(typeof(DataTestUtility), nameof(DataTestUtility.AreConnStringsSetup), nameof(DataTestUtility.IsNotAzureServer), nameof(DataTestUtility.IsSupportedDataClassification))]
18+
public static void TestDataClassificationResultSet()
19+
{
20+
s_tableName = DataTestUtility.GetUniqueNameForSqlServer("DC");
21+
using (SqlConnection sqlConnection = new SqlConnection(DataTestUtility.TCPConnectionString))
22+
using (SqlCommand sqlCommand = sqlConnection.CreateCommand())
23+
{
24+
try
25+
{
26+
sqlConnection.Open();
27+
Assert.True(DataTestUtility.IsSupportedDataClassification());
28+
CreateTable(sqlCommand);
29+
RunTestsForServer(sqlCommand);
30+
}
31+
finally
32+
{
33+
DataTestUtility.DropTable(sqlConnection, s_tableName);
34+
}
35+
}
36+
}
37+
38+
private static void RunTestsForServer(SqlCommand sqlCommand)
39+
{
40+
sqlCommand.CommandText = "SELECT * FROM " + s_tableName;
41+
using (SqlDataReader reader = sqlCommand.ExecuteReader())
42+
{
43+
VerifySensitivityClassification(reader);
44+
}
45+
}
46+
47+
private static void VerifySensitivityClassification(SqlDataReader reader)
48+
{
49+
if (null != reader.SensitivityClassification)
50+
{
51+
for (int columnPos = 0; columnPos < reader.SensitivityClassification.ColumnSensitivities.Count;
52+
columnPos++)
53+
{
54+
foreach (SensitivityProperty sp in reader.SensitivityClassification.ColumnSensitivities[columnPos].SensitivityProperties)
55+
{
56+
ReadOnlyCollection<InformationType> infoTypes = reader.SensitivityClassification.InformationTypes;
57+
Assert.Equal(3, infoTypes.Count);
58+
for (int i = 0; i < infoTypes.Count; i++)
59+
{
60+
VerifyInfoType(infoTypes[i], i + 1);
61+
}
62+
63+
ReadOnlyCollection<Label> labels = reader.SensitivityClassification.Labels;
64+
Assert.Single(labels);
65+
VerifyLabel(labels[0]);
66+
67+
if (columnPos == 1 || columnPos == 2 || columnPos == 6 || columnPos == 7)
68+
{
69+
VerifyLabel(sp.Label);
70+
VerifyInfoType(sp.InformationType, columnPos);
71+
}
72+
}
73+
}
74+
}
75+
}
76+
77+
private static void VerifyLabel(Label label)
78+
{
79+
Assert.True(label != null);
80+
Assert.Equal("L1", label.Id);
81+
Assert.Equal("PII", label.Name);
82+
}
83+
84+
private static void VerifyInfoType(InformationType informationType, int i)
85+
{
86+
Assert.True(informationType != null);
87+
Assert.Equal(i == 1 ? "COMPANY" : (i == 2 ? "NAME" : "CONTACT"), informationType.Id);
88+
Assert.Equal(i == 1 ? "Company Name" : (i == 2 ? "Person Name" : "Contact Information"), informationType.Name);
89+
}
90+
91+
private static void CreateTable(SqlCommand sqlCommand)
92+
{
93+
sqlCommand.CommandText = "CREATE TABLE " + s_tableName + " ("
94+
+ "[Id] [int] IDENTITY(1,1) NOT NULL,"
95+
+ "[CompanyName] [nvarchar](40) NOT NULL,"
96+
+ "[ContactName] [nvarchar](50) NULL,"
97+
+ "[ContactTitle] [nvarchar](40) NULL,"
98+
+ "[City] [nvarchar](40) NULL,"
99+
+ "[CountryName] [nvarchar](40) NULL,"
100+
+ "[Phone] [nvarchar](30) MASKED WITH (FUNCTION = 'default()') NULL,"
101+
+ "[Fax] [nvarchar](30) MASKED WITH (FUNCTION = 'default()') NULL)";
102+
sqlCommand.ExecuteNonQuery();
103+
104+
sqlCommand.CommandText = "ADD SENSITIVITY CLASSIFICATION TO " + s_tableName
105+
+ ".CompanyName WITH (LABEL='PII', LABEL_ID='L1', INFORMATION_TYPE='Company Name', INFORMATION_TYPE_ID='COMPANY')";
106+
sqlCommand.ExecuteNonQuery();
107+
108+
sqlCommand.CommandText = "ADD SENSITIVITY CLASSIFICATION TO " + s_tableName
109+
+ ".ContactName WITH (LABEL='PII', LABEL_ID='L1', INFORMATION_TYPE='Person Name', INFORMATION_TYPE_ID='NAME')";
110+
sqlCommand.ExecuteNonQuery();
111+
112+
sqlCommand.CommandText = "ADD SENSITIVITY CLASSIFICATION TO " + s_tableName
113+
+ ".Phone WITH (LABEL='PII', LABEL_ID='L1', INFORMATION_TYPE='Contact Information', INFORMATION_TYPE_ID='CONTACT')";
114+
sqlCommand.ExecuteNonQuery();
115+
116+
sqlCommand.CommandText = "ADD SENSITIVITY CLASSIFICATION TO " + s_tableName
117+
+ ".Fax WITH (LABEL='PII', LABEL_ID='L1', INFORMATION_TYPE='Contact Information', INFORMATION_TYPE_ID='CONTACT')";
118+
sqlCommand.ExecuteNonQuery();
119+
120+
// INSERT ROWS OF DATA
121+
sqlCommand.CommandText = "INSERT INTO " + s_tableName + " VALUES (@companyName, @contactName, @contactTitle, @city, @country, @phone, @fax)";
122+
123+
sqlCommand.Parameters.AddWithValue("@companyName", "Exotic Liquids");
124+
sqlCommand.Parameters.AddWithValue("@contactName", "Charlotte Cooper");
125+
sqlCommand.Parameters.AddWithValue("@contactTitle", "");
126+
sqlCommand.Parameters.AddWithValue("city", "London");
127+
sqlCommand.Parameters.AddWithValue("@country", "UK");
128+
sqlCommand.Parameters.AddWithValue("@phone", "(171) 555-2222");
129+
sqlCommand.Parameters.AddWithValue("@fax", "");
130+
sqlCommand.ExecuteNonQuery();
131+
132+
sqlCommand.Parameters.Clear();
133+
sqlCommand.Parameters.AddWithValue("@companyName", "New Orleans");
134+
sqlCommand.Parameters.AddWithValue("@contactName", "Cajun Delights");
135+
sqlCommand.Parameters.AddWithValue("@contactTitle", "");
136+
sqlCommand.Parameters.AddWithValue("city", "New Orleans");
137+
sqlCommand.Parameters.AddWithValue("@country", "USA");
138+
sqlCommand.Parameters.AddWithValue("@phone", "(100) 555-4822");
139+
sqlCommand.Parameters.AddWithValue("@fax", "");
140+
sqlCommand.ExecuteNonQuery();
141+
142+
sqlCommand.Parameters.Clear();
143+
sqlCommand.Parameters.AddWithValue("@companyName", "Grandma Kelly's Homestead");
144+
sqlCommand.Parameters.AddWithValue("@contactName", "Regina Murphy");
145+
sqlCommand.Parameters.AddWithValue("@contactTitle", "");
146+
sqlCommand.Parameters.AddWithValue("@city", "Ann Arbor");
147+
sqlCommand.Parameters.AddWithValue("@country", "USA");
148+
sqlCommand.Parameters.AddWithValue("@phone", "(313) 555-5735");
149+
sqlCommand.Parameters.AddWithValue("@fax", "(313) 555-3349");
150+
sqlCommand.ExecuteNonQuery();
151+
}
152+
153+
[ConditionalFact(typeof(DataTestUtility), nameof(DataTestUtility.AreConnStringsSetup), nameof(DataTestUtility.IsNotAzureServer), nameof(DataTestUtility.IsSupportedDataClassification))]
154+
public static void TestDataClassificationBulkCopy()
155+
{
156+
var data = new DataTable("Company");
157+
data.Columns.Add("CompanyId", typeof(Guid));
158+
data.Columns.Add("CompanyName", typeof(string));
159+
data.Columns.Add("Email", typeof(string));
160+
data.Columns.Add("CompanyType", typeof(int));
161+
162+
data.Rows.Add(Guid.NewGuid(), "Company 1", "[email protected]", 1);
163+
data.Rows.Add(Guid.NewGuid(), "Company 2", "[email protected]", 1);
164+
data.Rows.Add(Guid.NewGuid(), "Company 3", "[email protected]", 1);
165+
166+
var tableName = DataTestUtility.GetUniqueNameForSqlServer("DC");
167+
168+
using (var connection = new SqlConnection(DataTestUtility.TCPConnectionString))
169+
{
170+
connection.Open();
171+
try
172+
{
173+
// Setup Table
174+
using (SqlCommand sqlCommand = connection.CreateCommand())
175+
{
176+
sqlCommand.CommandText = $"CREATE TABLE {tableName} (" +
177+
$" [CompanyId] [uniqueidentifier] NOT NULL," +
178+
$" [CompanyName][nvarchar](255) NOT NULL," +
179+
$" [Email] [nvarchar](50) NULL," +
180+
$" [CompanyType] [int] not null," +
181+
$" CONSTRAINT[PK_Company] PRIMARY KEY CLUSTERED (" +
182+
$" [CompanyId] ASC" +
183+
$" ) WITH(STATISTICS_NORECOMPUTE = OFF, IGNORE_DUP_KEY = OFF) ON[PRIMARY]" +
184+
$" ) ON[PRIMARY]";
185+
sqlCommand.ExecuteNonQuery();
186+
187+
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');";
188+
sqlCommand.ExecuteNonQuery();
189+
}
190+
191+
// Perform Bulk Insert
192+
using (var bulk = new SqlBulkCopy(connection))
193+
{
194+
bulk.DestinationTableName = tableName;
195+
bulk.ColumnMappings.Add("CompanyId", "CompanyId");
196+
bulk.ColumnMappings.Add("CompanyName", "CompanyName");
197+
bulk.ColumnMappings.Add("Email", "Email");
198+
bulk.ColumnMappings.Add("CompanyType", "CompanyType");
199+
bulk.WriteToServer(data);
200+
}
201+
}
202+
finally
203+
{
204+
DataTestUtility.DropTable(connection, tableName);
205+
}
206+
}
207+
}
208+
}
209+
}

0 commit comments

Comments
 (0)