Handle preview mode for execute endpoint

InseeFrLab · Oct 9, 2024 · bb97008 · bb97008
1 parent ee209dc
commit bb97008
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 8 deletions.
diff --git a/src/main/java/fr/insee/trevas/lab/controller/TrevasLabController.java b/src/main/java/fr/insee/trevas/lab/controller/TrevasLabController.java
@@ -81,13 +81,14 @@ public ResponseEntity<String> postProvenance(
     public ResponseEntity<UUID> executeNew(
             Authentication auth,
             @RequestBody Body body,
-            @RequestParam("mode") ExecutionMode mode
+            @RequestParam("mode") ExecutionMode mode,
+            @RequestParam("preview") Boolean preview
     ) throws Exception {
         Job job;
         if (mode == ExecutionMode.MEMORY) {
             job = executeJob(body, () -> {
                 try {
-                    return inMemoryEngine.executeInMemory(userProvider.getUser(auth), body);
+                    return inMemoryEngine.executeInMemory(userProvider.getUser(auth), body, preview);
                 } catch (Exception e) {
                     e.printStackTrace();
                     throw new Exception(e.getMessage());
@@ -96,7 +97,7 @@ public ResponseEntity<UUID> executeNew(
         } else if (mode == ExecutionMode.SPARK) {
             job = executeJob(body, () -> {
                 try {
-                    return sparkEngine.executeSpark(userProvider.getUser(auth), body);
+                    return sparkEngine.executeSpark(userProvider.getUser(auth), body, preview);
                 } catch (Exception e) {
                     e.printStackTrace();
                     throw new Exception(e.getMessage());

diff --git a/src/main/java/fr/insee/trevas/lab/service/InMemoryEngine.java b/src/main/java/fr/insee/trevas/lab/service/InMemoryEngine.java
@@ -30,7 +30,7 @@ public class InMemoryEngine {
     @Autowired
     private ObjectMapper objectMapper;
 
-    public Bindings executeInMemory(User user, Body body) throws Exception {
+    public Bindings executeInMemory(User user, Body body, Boolean preview) throws Exception {
         String script = body.getVtlScript();
         Bindings bindings = body.getBindings();
         Map<String, QueriesForBindings> queriesForBindings = body.getQueriesForBindings();
@@ -53,7 +53,8 @@ public Bindings executeInMemory(User user, Body body) throws Exception {
                                     v.getPassword())
                     ) {
                         Statement statement = connection.createStatement();
-                        return statement.executeQuery(v.getQuery());
+                        String query = preview ? v.getQuery() + " LIMIT 0" : v.getQuery();
+                        return statement.executeQuery(query);
                     } catch (SQLException se) {
                         throw new RuntimeException(se);
                     }

diff --git a/src/main/java/fr/insee/trevas/lab/service/SparkEngine.java b/src/main/java/fr/insee/trevas/lab/service/SparkEngine.java
@@ -107,7 +107,7 @@ private SparkDataset readJDBCDataset(SparkSession spark, QueriesForBindings quer
         return new SparkDataset(dataset);
     }
 
-    public Bindings executeSpark(User user, Body body) throws Exception {
+    public Bindings executeSpark(User user, Body body, Boolean preview) throws Exception {
         String script = body.getVtlScript();
         Map<String, QueriesForBindings> queriesForBindings = body.getQueriesForBindings();
         Map<String, S3ForBindings> s3ForBindings = body.getS3ForBindings();
@@ -116,10 +116,13 @@ public Bindings executeSpark(User user, Body body) throws Exception {
 
         Bindings bindings = new SimpleBindings();
 
+        Integer limit = preview ? 0 : null;
+
         if (queriesForBindings != null) {
             queriesForBindings.forEach((k, v) -> {
                 try {
-                    SparkDataset sparkDataset = readJDBCDataset(spark, v, null);
+
+                    SparkDataset sparkDataset = readJDBCDataset(spark, v, limit);
                     bindings.put(k, sparkDataset);
                 } catch (Exception e) {
                     logger.warn("Query loading failed: ", e);
@@ -129,7 +132,7 @@ public Bindings executeSpark(User user, Body body) throws Exception {
         if (s3ForBindings != null) {
             s3ForBindings.forEach((k, v) -> {
                 try {
-                    SparkDataset sparkDataset = readS3Dataset(spark, v, null);
+                    SparkDataset sparkDataset = readS3Dataset(spark, v, limit);
                     bindings.put(k, sparkDataset);
                 } catch (Exception e) {
                     logger.warn("S3 loading failed: ", e);