Query execution

Author: Fabricio20

build.gradle

@@ -27,6 +27,7 @@ dependencies {
     implementation('org.springframework.boot:spring-boot-starter-security')
 
     implementation('org.jetbrains:annotations:23.0.0')
+    implementation('com.opencsv:opencsv:5.7.1')
 
     compileOnly 'org.projectlombok:lombok'
     runtimeOnly 'org.postgresql:postgresql'

src/main/java/com/scriptql/api/domain/entities/Query.java

@@ -38,6 +38,9 @@ public class Query extends BaseEntity {
     @JsonIgnore
     private byte[] result;
 
+    @Nullable
+    private String error;
+
     @Nullable
     @JsonProperty("execution_date")
     @Column(name = "execution_date")

src/main/java/com/scriptql/api/domain/enums/QueryStatus.java

@@ -1,9 +1,12 @@
 package com.scriptql.api.domain.enums;
 
 public enum QueryStatus {
+
     WAITING_REVIEW,
     APPROVED,
     REJECTED,
     EXECUTING,
-    DONE
+    DONE,
+    ERROR
+
 }

src/main/java/com/scriptql/api/rest/controllers/QueryController.java

@@ -7,6 +7,7 @@
 import com.scriptql.api.domain.request.CreateQueryRequest;
 import com.scriptql.api.services.QueryService;
 import org.springframework.http.HttpStatus;
+import org.springframework.http.ResponseEntity;
 import org.springframework.web.bind.annotation.*;
 
 import javax.validation.Valid;
@@ -43,4 +44,14 @@ public List<Review> getReviews(@PathVariable("id") long id) {
         return this.service.getReviews(id);
     }
 
+    @PostMapping("/{id:\\d+}/execute")
+    public void execute(@PathVariable("id") long id) {
+        this.service.execute(id);
+    }
+
+    @GetMapping("/{id:\\d+}/download")
+    public ResponseEntity<byte[]> download(@PathVariable("id") long id) {
+        return this.service.download(id);
+    }
+
 }

src/main/java/com/scriptql/api/services/ExecutionService.java

@@ -0,0 +1,80 @@
+package com.scriptql.api.services;
+
+import com.opencsv.CSVWriter;
+import com.scriptql.api.domain.entities.DatabaseConnection;
+import com.scriptql.api.domain.entities.Query;
+import com.scriptql.api.domain.enums.DatabaseDriver;
+import com.scriptql.api.domain.enums.QueryStatus;
+import com.scriptql.api.domain.repositories.QueryRepository;
+import org.springframework.stereotype.Service;
+
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.OutputStreamWriter;
+import java.sql.*;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+
+@Service
+public class ExecutionService {
+
+    private final QueryRepository repository;
+    private final ExecutorService scheduler = Executors.newCachedThreadPool();
+
+    public ExecutionService(QueryRepository repository) {
+        this.repository = repository;
+    }
+
+    public void execute(Query query) {
+        this.scheduler.submit(() -> this.onExecute(query));
+    }
+
+    private void onExecute(Query query) {
+        try (Connection connection = this.open(query.getConnection())) {
+            String sql = query.getQuery();
+            boolean isSelect = sql.toUpperCase().startsWith("SELECT");
+            if (isSelect) {
+                connection.setReadOnly(true);
+            }
+            try (PreparedStatement ps = connection.prepareStatement(sql)) {
+                if (isSelect) {
+                    try (ResultSet rs = ps.executeQuery()) {
+                        ByteArrayOutputStream stream = new ByteArrayOutputStream();
+                        try (CSVWriter writer = new CSVWriter(new OutputStreamWriter(stream))) {
+                            writer.writeAll(rs, true, false, true);
+                        }
+                        query.setResult(stream.toByteArray());
+                    }
+                } else {
+                    int result = ps.executeUpdate();
+                    ByteArrayOutputStream stream = new ByteArrayOutputStream();
+                    try (CSVWriter writer = new CSVWriter(new OutputStreamWriter(stream))) {
+                        writer.writeNext(new String[]{"updated", String.valueOf(result)});
+                    }
+                    query.setResult(stream.toByteArray());
+                }
+            }
+            query.setStatus(QueryStatus.DONE);
+        } catch (SQLException | IOException e) {
+            query.setStatus(QueryStatus.ERROR);
+            query.setError(e.getMessage());
+        } finally {
+            query.setExecutionDate(System.currentTimeMillis());
+            query.setUpdatedAt(System.currentTimeMillis());
+            this.repository.save(query);
+        }
+    }
+
+    private Connection open(DatabaseConnection details) throws SQLException {
+        String url;
+        if (details.getDriver() == DatabaseDriver.MYSQL) {
+            url = String.format("jdbc:mysql://%s:%s/%s",
+                    details.getHost(), details.getPort(), details.getDatabase());
+        } else {
+            url = String.format("jdbc:postgresql://%s:%s/%s",
+                    details.getHost(), details.getPort(), details.getDatabase());
+        }
+        return DriverManager.getConnection(url, details.getUsername(), details.getPassword());
+    }
+
+}

src/main/java/com/scriptql/api/services/QueryService.java

@@ -15,8 +15,10 @@
 import com.scriptql.api.domain.request.CreateQueryRequest;
 import org.jetbrains.annotations.NotNull;
 import org.springframework.data.jpa.domain.Specification;
+import org.springframework.http.*;
 import org.springframework.stereotype.Service;
 
+import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Random;
@@ -30,6 +32,7 @@ public class QueryService {
     private final RoleRepository roles;
     private final ReviewRepository reviews;
     private final NotificationService notification;
+    private final ExecutionService executor;
 
     private final SpecBuilder<Query> builder = new SpecBuilder<>();
 
@@ -38,13 +41,14 @@ public QueryService(
             ConnectionService connections,
             UserRoleRepository userRoles,
             RoleRepository roles, ReviewRepository reviews,
-            NotificationService notification) {
+            NotificationService notification, ExecutionService executor) {
         this.repository = repository;
         this.connections = connections;
         this.userRoles = userRoles;
         this.roles = roles;
         this.reviews = reviews;
         this.notification = notification;
+        this.executor = executor;
         this.builder.addMatcher("database.name", SpecMatcher.SEARCH);
         this.builder.addMatcher("description", SpecMatcher.SEARCH);
     }
@@ -97,4 +101,31 @@ public QueryService(
         return this.reviews.findAllByQuery(query);
     }
 
+    public void execute(long id) {
+        Query query = this.findById(id);
+        if (query.getStatus() != QueryStatus.APPROVED) {
+            throw new UserError("Unable to execute this query");
+        }
+        query.setStatus(QueryStatus.EXECUTING);
+        this.repository.save(query);
+        this.executor.execute(query);
+    }
+
+    public ResponseEntity<byte[]> download(long id) {
+        Query query = this.findById(id);
+        if (query.getStatus() != QueryStatus.DONE) {
+            throw new UserError("This query has not executed yet");
+        }
+        byte[] result = query.getResult();
+        if (result == null || result.length == 0) {
+            return ResponseEntity.noContent().build();
+        }
+        HttpHeaders headers = new HttpHeaders();
+        headers.setContentType(MediaType.parseMediaType("text/csv"));
+        headers.setContentDisposition(ContentDisposition.attachment()
+                .filename(id + ".csv", StandardCharsets.UTF_8)
+                .build());
+        return new ResponseEntity<>(result, headers, HttpStatus.OK);
+    }
+
 }