clean up

linkedin · Sep 14, 2023 · 08633fa · 08633fa
1 parent bd30503
commit 08633fa
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 17 deletions.
diff --git a/...ice/src/main/java/com/linkedin/coral/coralservice/controller/VisualizationController.java b/...ice/src/main/java/com/linkedin/coral/coralservice/controller/VisualizationController.java
@@ -62,7 +62,7 @@ public ResponseEntity getIRVisualizations(@RequestBody VisualizationRequestBody
       return ResponseEntity.status(HttpStatus.INTERNAL_SERVER_ERROR).body(t.getMessage());
     }
 
-    // Build  response body
+    // Build response body
     VisualizationResponseBody responseBody = new VisualizationResponseBody();
     responseBody.setSqlNodeImageID(sqlNodeImageID);
     responseBody.setRelNodeImageID(relNodeImageID);
@@ -76,7 +76,7 @@ public ResponseEntity getIRVisualizations(@RequestBody VisualizationRequestBody
   public ResponseEntity<FileSystemResource> getImage(@PathVariable String imageId) {
     String imagePath = imageDir + File.separator + imageId + ".svg";
 
-    if (isValidImageId(imagePath)) {
+    if (isValidImage(imagePath)) {
       try {
         Path path = new File(imagePath).toPath();
         String contentType = Files.probeContentType(path);
@@ -97,9 +97,9 @@ public ResponseEntity<FileSystemResource> getImage(@PathVariable String imageId)
     }
   }
 
-  private boolean isValidImageId(String imagePath) {
-    File imageFile = new File(imagePath);
+  private boolean isValidImage(String imagePath) {
     // Check if the file exists and is a regular file (not a directory)
+    File imageFile = new File(imagePath);
     return imageFile.exists() && imageFile.isFile();
   }
 }
diff --git a/coral-service/src/main/java/com/linkedin/coral/coralservice/utils/VisualizationUtils.java b/coral-service/src/main/java/com/linkedin/coral/coralservice/utils/VisualizationUtils.java
@@ -39,22 +39,15 @@ public static UUID generateSqlNodeVisualization(String query, String fromLanguag
       sqlNode = new HiveToRelConverter(hiveMetastoreClient).toSqlNode(query);
     }
 
-    if (incrementalRewrittenRelNode != null) {
-      switch (rewriteType) {
-        case INCREMENTAL:
-          // We want to instead generate the visualization of SqlNode2 of the RHS of Coral's translation
-          sqlNode = new CoralRelToSqlNodeConverter().convert(incrementalRewrittenRelNode);
-          break;
-        case DATAMASKING:
-        case NONE:
-        default:
-          break;
-      }
+    if (incrementalRewrittenRelNode != null && rewriteType == INCREMENTAL) {
+      // We want to instead generate the visualization of SqlNode2 of the RHS of Coral's translation
+      sqlNode = new CoralRelToSqlNodeConverter().convert(incrementalRewrittenRelNode);
     }
 
     assert sqlNode != null;
-    VisualizationUtil visualizationUtil = VisualizationUtil.create(imageDir);
 
+    // Generate graphviz svg using sqlNode
+    VisualizationUtil visualizationUtil = VisualizationUtil.create(imageDir);
     UUID sqlNodeId = UUID.randomUUID();
     visualizationUtil.visualizeSqlNodeToFile(sqlNode, "/" + sqlNodeId + ".svg");
 
@@ -82,8 +75,9 @@ public static UUID generateRelNodeVisualization(String query, String fromLanguag
     }
 
     assert relNode != null;
-    VisualizationUtil visualizationUtil = VisualizationUtil.create(imageDir);
 
+    // Generate graphviz svg using relNode
+    VisualizationUtil visualizationUtil = VisualizationUtil.create(imageDir);
     UUID relNodeID = UUID.randomUUID();
     visualizationUtil.visualizeRelNodeToFile(relNode, "/" + relNodeID + ".svg");