-
Notifications
You must be signed in to change notification settings - Fork 188
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[Coral-Service] Graph visualization API (with rewrites) (#454)
- Loading branch information
Showing
8 changed files
with
358 additions
and
17 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
111 changes: 111 additions & 0 deletions
111
...ice/src/main/java/com/linkedin/coral/coralservice/controller/VisualizationController.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,111 @@ | ||
| /** | ||
| * Copyright 2023 LinkedIn Corporation. All rights reserved. | ||
| * Licensed under the BSD-2 Clause license. | ||
| * See LICENSE in the project root for license information. | ||
| */ | ||
| package com.linkedin.coral.coralservice.controller; | ||
|
|
||
| import java.io.File; | ||
| import java.io.IOException; | ||
| import java.nio.file.Files; | ||
| import java.nio.file.Path; | ||
| import java.util.ArrayList; | ||
| import java.util.UUID; | ||
|
|
||
| import org.springframework.core.io.FileSystemResource; | ||
| import org.springframework.http.HttpStatus; | ||
| import org.springframework.http.MediaType; | ||
| import org.springframework.http.ResponseEntity; | ||
| import org.springframework.web.bind.annotation.GetMapping; | ||
| import org.springframework.web.bind.annotation.PathVariable; | ||
| import org.springframework.web.bind.annotation.PostMapping; | ||
| import org.springframework.web.bind.annotation.RequestBody; | ||
| import org.springframework.web.bind.annotation.RequestMapping; | ||
| import org.springframework.web.bind.annotation.RestController; | ||
|
|
||
| import com.linkedin.coral.coralservice.entity.VisualizationRequestBody; | ||
| import com.linkedin.coral.coralservice.entity.VisualizationResponseBody; | ||
| import com.linkedin.coral.coralservice.utils.RewriteType; | ||
| import com.linkedin.coral.coralservice.utils.VisualizationUtils; | ||
|
|
||
| import static com.linkedin.coral.coralservice.utils.VisualizationUtils.*; | ||
|
|
||
|
|
||
| @RestController | ||
| @RequestMapping("/api/visualizations") | ||
| public class VisualizationController { | ||
| private File imageDir = getImageDir(); | ||
| private VisualizationUtils visualizationUtils = new VisualizationUtils(); | ||
|
|
||
| @PostMapping("/generategraphs") | ||
| public ResponseEntity getIRVisualizations(@RequestBody VisualizationRequestBody visualizationRequestBody) { | ||
| final String sourceLanguage = visualizationRequestBody.getSourceLanguage(); | ||
| final String query = visualizationRequestBody.getQuery(); | ||
| final RewriteType rewriteType = visualizationRequestBody.getRewriteType(); | ||
|
|
||
| if (!visualizationUtils.isValidSourceLanguage(sourceLanguage)) { | ||
| return ResponseEntity.status(HttpStatus.BAD_REQUEST) | ||
| .body("Currently, only Hive, Spark, and Trino are supported as engines to generate graphs using.\n"); | ||
| } | ||
|
|
||
| // A list of UUIDs in this order of: | ||
| // 1. Image ID of pre/no rewrite relNode | ||
| // 2. Image ID of pre/no rewrite sqlNode | ||
| // If a rewrite was requested: | ||
| // 3. Image ID of post rewrite relNode | ||
| // 4. Image ID of post rewrite sqlNode | ||
| ArrayList<UUID> imageIdList; | ||
| try { | ||
| imageIdList = visualizationUtils.generateIRVisualizations(query, sourceLanguage, imageDir, rewriteType); | ||
| } catch (Throwable t) { | ||
| // TODO: use logger | ||
| t.printStackTrace(); | ||
| return ResponseEntity.status(HttpStatus.INTERNAL_SERVER_ERROR).body(t.getMessage()); | ||
| } | ||
|
|
||
| assert imageIdList.size() > 0; | ||
|
|
||
| // Build response body | ||
| VisualizationResponseBody responseBody = new VisualizationResponseBody(); | ||
| responseBody.setRelNodeImageID(imageIdList.get(0)); | ||
| responseBody.setSqlNodeImageID(imageIdList.get(1)); | ||
| if (imageIdList.size() >= 4) { | ||
| // Rewrite was requested | ||
| responseBody.setPostRewriteRelNodeImageID(imageIdList.get(2)); | ||
| responseBody.setPostRewriteSqlNodeImageID(imageIdList.get(3)); | ||
| } | ||
|
|
||
| return ResponseEntity.status(HttpStatus.OK).body(responseBody); | ||
| } | ||
|
|
||
| @GetMapping("/{imageId}") | ||
| public ResponseEntity<FileSystemResource> getImage(@PathVariable String imageId) { | ||
| String imagePath = imageDir + File.separator + imageId + ".svg"; | ||
|
|
||
| if (isValidImage(imagePath)) { | ||
| try { | ||
| Path path = new File(imagePath).toPath(); | ||
| String contentType = Files.probeContentType(path); | ||
|
|
||
| if (contentType == null) { | ||
| contentType = "image/svg+xml"; | ||
| } | ||
| FileSystemResource resource = new FileSystemResource(path); | ||
| return ResponseEntity.ok().contentType(MediaType.parseMediaType(contentType)).body(resource); | ||
|
|
||
| } catch (IOException e) { | ||
| e.printStackTrace(); | ||
| return ResponseEntity.status(500).build(); // 500 Internal Server Error response | ||
| } | ||
|
|
||
| } else { | ||
| return ResponseEntity.notFound().build(); | ||
| } | ||
| } | ||
|
|
||
| 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(); | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
28 changes: 28 additions & 0 deletions
28
...ervice/src/main/java/com/linkedin/coral/coralservice/entity/VisualizationRequestBody.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,28 @@ | ||
| /** | ||
| * Copyright 2022-2023 LinkedIn Corporation. All rights reserved. | ||
| * Licensed under the BSD-2 Clause license. | ||
| * See LICENSE in the project root for license information. | ||
| */ | ||
| package com.linkedin.coral.coralservice.entity; | ||
|
|
||
| import com.linkedin.coral.coralservice.utils.RewriteType; | ||
|
|
||
|
|
||
| public class VisualizationRequestBody { | ||
| private String sourceLanguage; | ||
| private String query; | ||
|
|
||
| private RewriteType rewriteType; | ||
|
|
||
| public String getSourceLanguage() { | ||
| return sourceLanguage; | ||
| } | ||
|
|
||
| public String getQuery() { | ||
| return query; | ||
| } | ||
|
|
||
| public RewriteType getRewriteType() { | ||
| return rewriteType; | ||
| } | ||
| } |
52 changes: 52 additions & 0 deletions
52
...rvice/src/main/java/com/linkedin/coral/coralservice/entity/VisualizationResponseBody.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,52 @@ | ||
| /** | ||
| * Copyright 2023 LinkedIn Corporation. All rights reserved. | ||
| * Licensed under the BSD-2 Clause license. | ||
| * See LICENSE in the project root for license information. | ||
| */ | ||
| package com.linkedin.coral.coralservice.entity; | ||
|
|
||
| import java.util.UUID; | ||
|
|
||
|
|
||
| public class VisualizationResponseBody { | ||
|
|
||
| private UUID sqlNodeImageID; | ||
| private UUID relNodeImageID; | ||
| private UUID postRewriteSqlNodeImageID; | ||
| private UUID postRewriteRelNodeImageID; | ||
| public VisualizationResponseBody() { | ||
| } | ||
|
|
||
| public UUID getSqlNodeImageID() { | ||
| return sqlNodeImageID; | ||
| } | ||
|
|
||
| public UUID getRelNodeImageID() { | ||
| return relNodeImageID; | ||
| } | ||
|
|
||
| public UUID getPostRewriteSqlNodeImageID() { | ||
| return postRewriteSqlNodeImageID; | ||
| } | ||
|
|
||
| public void setPostRewriteSqlNodeImageID(UUID postRewriteSqlNodeImageID) { | ||
| this.postRewriteSqlNodeImageID = postRewriteSqlNodeImageID; | ||
| } | ||
|
|
||
| public void setSqlNodeImageID(UUID sqlNodeImageID) { | ||
| this.sqlNodeImageID = sqlNodeImageID; | ||
| } | ||
|
|
||
| public void setRelNodeImageID(UUID relNodeImageID) { | ||
| this.relNodeImageID = relNodeImageID; | ||
| } | ||
|
|
||
| public UUID getPostRewriteRelNodeImageID() { | ||
| return postRewriteRelNodeImageID; | ||
| } | ||
|
|
||
| public void setPostRewriteRelNodeImageID(UUID postRewriteRelNodeImageID) { | ||
| this.postRewriteRelNodeImageID = postRewriteRelNodeImageID; | ||
| } | ||
|
|
||
| } |
37 changes: 37 additions & 0 deletions
37
coral-service/src/main/java/com/linkedin/coral/coralservice/utils/RewriteType.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,37 @@ | ||
| /** | ||
| * Copyright 2023 LinkedIn Corporation. All rights reserved. | ||
| * Licensed under the BSD-2 Clause license. | ||
| * See LICENSE in the project root for license information. | ||
| */ | ||
| package com.linkedin.coral.coralservice.utils; | ||
|
|
||
| import com.fasterxml.jackson.annotation.JsonCreator; | ||
|
|
||
|
|
||
| public enum RewriteType { | ||
| NONE("none"), | ||
| INCREMENTAL("incremental"), | ||
| DATAMASKING("datamasking"); | ||
|
|
||
| private final String type; | ||
|
|
||
| RewriteType(String type) { | ||
| this.type = type; | ||
| } | ||
|
|
||
| @Override | ||
| public String toString() { | ||
| return type; | ||
| } | ||
|
|
||
| @JsonCreator | ||
| public static RewriteType getDepartmentFromCode(String value) { | ||
| for (RewriteType type : RewriteType.values()) { | ||
| if (type.toString().equals(value)) { | ||
| return type; | ||
| } | ||
| } | ||
| return null; | ||
| } | ||
|
|
||
| } |
Oops, something went wrong.