Dataclassification sensitivity max index fix (#1847)

microsoft · Jun 22, 2022 · c9f955f · c9f955f
1 parent a417222
commit c9f955f
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 2 deletions.
diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/StreamColumns.java b/src/main/java/com/microsoft/sqlserver/jdbc/StreamColumns.java
@@ -285,7 +285,9 @@ SensitivityClassification processDataClassification(TDSReader tdsReader) throws
                 // get the label index and then lookup label to use for source
                 int sensitivityLabelIndex = tdsReader.readUnsignedShort();
                 Label label = null;
-                if (sensitivityLabelIndex != Integer.MAX_VALUE) {
+
+                // If sensitivity label index is equal to USHORT_MAX eg. 65535 then there is no sensitivity label
+                if (sensitivityLabelIndex != DataTypes.SQL_USHORTVARMAXLEN) {
                     if (sensitivityLabelIndex >= sensitivityLabels.size()) {
                         tdsReader.throwInvalidTDS();
                     }
@@ -294,7 +296,9 @@ SensitivityClassification processDataClassification(TDSReader tdsReader) throws
                 // get the information type index and then lookup information type to use for source
                 int informationTypeIndex = tdsReader.readUnsignedShort();
                 InformationType informationType = null;
-                if (informationTypeIndex != Integer.MAX_VALUE) {
+
+                // If information type index is equal to USHORT_MAX eg. 65535 then there is no information type
+                if (informationTypeIndex != DataTypes.SQL_USHORTVARMAXLEN) {
                     if (informationTypeIndex >= informationTypes.size()) {}
                     informationType = informationTypes.get(informationTypeIndex);
                 }

diff --git a/src/test/java/com/microsoft/sqlserver/jdbc/resultset/DataClassificationTest.java b/src/test/java/com/microsoft/sqlserver/jdbc/resultset/DataClassificationTest.java
@@ -44,6 +44,8 @@ public class DataClassificationTest extends AbstractTest {
             .escapeIdentifier(RandomUtil.getIdentifier("DataClassification"));
     private static final String tableName1 = AbstractSQLGenerator
             .escapeIdentifier(RandomUtil.getIdentifier("SelectMethodCursor"));
+    private static final String tableName2 = AbstractSQLGenerator
+            .escapeIdentifier(RandomUtil.getIdentifier("USHORT_MAX"));
     private static final String addSensitivitySql = "ADD SENSITIVITY CLASSIFICATION TO %s.%s WITH (LABEL='PII', LABEL_ID='L1', INFORMATION_TYPE='%s', INFORMATION_TYPE_ID='%s'%s)";
     private static final String sensitivityRankSql = ", RANK=%s";
 
@@ -52,6 +54,34 @@ public static void setupTests() throws Exception {
         setConnection();
     }
 
+    @Test
+    @Tag(Constants.xAzureSQLDW)
+    @Tag(Constants.xSQLv11)
+    @Tag(Constants.xSQLv12)
+    @Tag(Constants.xSQLv14)
+    public void testInformationTypeSensitivityLabelIndexMax() throws Exception {
+        String createTable = "create table " + tableName2 + " (col1 varchar(200), col2 varchar(200), col3 varchar(200))";
+        String sensitivityClassification = "ADD SENSITIVITY CLASSIFICATION TO %s.%s WITH (LABEL='PII')";
+        String query = "select * from " + tableName2;
+
+        try (SQLServerConnection conn = PrepUtil.getConnection(connectionString)) {
+            // Create table
+            try (Statement stmt = conn.createStatement()) {
+                stmt.execute(createTable);
+            }
+
+            // Add sensitivity classification without an INFORMATION_TYPE
+            try (Statement stmt = conn.createStatement()) {
+                stmt.execute(String.format(sensitivityClassification, tableName2, "col1"));
+            }
+
+            // Execute query
+            try (Statement stmt = conn.createStatement()) {
+                stmt.execute(query);
+            }
+        }
+    }
+
     @Test
     @Tag(Constants.xAzureSQLDW)
     @Tag(Constants.xSQLv11)