From ab31c4706e4e0af98badecaf9db4100c09a40a22 Mon Sep 17 00:00:00 2001
From: "Yi (Alan) Wang" <ywang5@linkedin.com>
Date: Wed, 16 Aug 2017 08:59:22 -0700
Subject: [PATCH] Update Dataset compliance DAO (#674)

---
 wherehows-dao/src/main/java/wherehows/dao/DatasetsDao.java  | 6 +++++-
 .../src/main/java/wherehows/models/DatasetCompliance.java   | 6 +++---
 2 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/wherehows-dao/src/main/java/wherehows/dao/DatasetsDao.java b/wherehows-dao/src/main/java/wherehows/dao/DatasetsDao.java
index 499e9fb7f2..0c8d2e67c4 100644
--- a/wherehows-dao/src/main/java/wherehows/dao/DatasetsDao.java
+++ b/wherehows-dao/src/main/java/wherehows/dao/DatasetsDao.java
@@ -139,7 +139,7 @@ public class DatasetsDao {
     return record;
   }
 
-  public void updateDatasetComplianceInfo(JdbcTemplate jdbcTemplate,
+  public DatasetCompliance updateDatasetComplianceInfo(JdbcTemplate jdbcTemplate,
       NamedParameterJdbcTemplate namedParameterJdbcTemplate, int datasetId, JsonNode node, String user)
       throws Exception {
 
@@ -148,8 +148,10 @@ public class DatasetsDao {
     if (record.getDatasetId() != null && datasetId != record.getDatasetId()) {
       throw new IllegalArgumentException("Dataset id doesn't match.");
     }
+    record.setDatasetId(datasetId);
 
     String urn = record.getDatasetUrn() != null ? record.getDatasetUrn() : getDatasetUrnById(jdbcTemplate, datasetId);
+    record.setDatasetUrn(urn);
 
     Map<String, Object> parameters = new HashMap<>();
     parameters.put("id", datasetId);
@@ -165,5 +167,7 @@ public class DatasetsDao {
     parameters.put("modified_by", user);
     parameters.put("modified_time", System.currentTimeMillis() / 1000);
     namedParameterJdbcTemplate.update(INSERT_DATASET_COMPLIANCE, parameters);
+
+    return record;
   }
 }
diff --git a/wherehows-dao/src/main/java/wherehows/models/DatasetCompliance.java b/wherehows-dao/src/main/java/wherehows/models/DatasetCompliance.java
index 72f6305290..4e599175e5 100644
--- a/wherehows-dao/src/main/java/wherehows/models/DatasetCompliance.java
+++ b/wherehows-dao/src/main/java/wherehows/models/DatasetCompliance.java
@@ -24,7 +24,7 @@ public class DatasetCompliance {
   private String complianceType;
   private List<DatasetFieldEntity> complianceEntities;
   private String confidentiality;
-  private Map<String, Object> datasetClassification;
+  private Map<String, Boolean> datasetClassification;
   private Map<String, String> fieldClassification;
   private String recordOwnerType;
   private Map<String, Object> retentionPolicy;
@@ -75,11 +75,11 @@ public class DatasetCompliance {
     this.confidentiality = confidentiality;
   }
 
-  public Map<String, Object> getDatasetClassification() {
+  public Map<String, Boolean> getDatasetClassification() {
     return datasetClassification;
   }
 
-  public void setDatasetClassification(Map<String, Object> datasetClassification) {
+  public void setDatasetClassification(Map<String, Boolean> datasetClassification) {
     this.datasetClassification = datasetClassification;
   }