Feat: EXPORT SYNOPSIS ~ TO FILE

traindb-project · Mar 27, 2024 · a18123f · a18123f
1 parent 7d45b90
commit a18123f
Show file tree

Hide file tree

Showing 5 changed files with 39 additions and 10 deletions.
diff --git a/traindb-core/src/main/antlr4/traindb/sql/TrainDBSql.g4 b/traindb-core/src/main/antlr4/traindb/sql/TrainDBSql.g4
@@ -158,7 +158,11 @@ modelBinaryString
     ;
 
 exportSynopsis
-    : K_EXPORT K_SYNOPSIS synopsisName
+    : K_EXPORT K_SYNOPSIS synopsisName exportToClause?
+    ;
+
+exportToClause
+    : K_TO K_FILE filenameString
     ;
 
 importSynopsis
@@ -169,6 +173,10 @@ analyzeSynopsis
     : K_ANALYZE K_SYNOPSIS synopsisName
     ;
 
+filenameString
+    : STRING_LITERAL
+    ;
+
 synopsisBinaryString
     : BINARY_STRING_LITERAL
     ;
@@ -318,6 +326,7 @@ K_DISABLE : D I S A B L E ;
 K_DROP : D R O P ;
 K_ENABLE : E N A B L E ;
 K_EXPORT : E X P O R T ;
+K_FILE : F I L E ;
 K_FOR : F O R ;
 K_FROM : F R O M ;
 K_HYPERPARAMETERS : H Y P E R P A R A M E T E R S ;

diff --git a/traindb-core/src/main/java/traindb/engine/TrainDBQueryEngine.java b/traindb-core/src/main/java/traindb/engine/TrainDBQueryEngine.java
@@ -1286,7 +1286,8 @@ public void disableModel(String modelName) throws Exception {
   }
 
   @Override
-  public TrainDBListResultSet exportSynopsis(String synopsisName) throws Exception {
+  public TrainDBListResultSet exportSynopsis(String synopsisName, String exportFilename)
+      throws Exception {
     T_tracer.startTaskTracer("export synopsis " + synopsisName);
 
     T_tracer.openTaskTime("find : synopsis");
@@ -1328,20 +1329,29 @@ public TrainDBListResultSet exportSynopsis(String synopsisName) throws Exception
       JdbcUtils.close(extConn, stmt, synopsisData);
     }
 
-    Path outputPath = Paths.get(tempDir.toString(), synopsisName + ".zip");
+    Path outputPath;
+    if (exportFilename != null) {
+      outputPath = Paths.get(exportFilename).toAbsolutePath();
+    } else {
+      outputPath = Paths.get(tempDir.toString(), synopsisName + ".zip");
+    }
     ZipUtils.pack(synopsisDir.toString(), outputPath.toString());
     ObjectMapper mapper = new ObjectMapper();
     ZipUtils.addNewFileFromStringToZip("export_synopsis.json",
         mapper.writeValueAsString(mSynopsis), outputPath);
 
+    T_tracer.closeTaskTime("SUCCESS");
+    T_tracer.endTaskTracer();
+
+    if (exportFilename != null) {
+      return TrainDBListResultSet.empty();
+    }
+
     List<String> header = Arrays.asList("export_synopsis");
     List<List<Object>> exportSynopsisInfo = new ArrayList<>();
     ByteArray byteArray = convertFileToByteArray(new File(outputPath.toString()));
     exportSynopsisInfo.add(Arrays.asList(byteArray));
 
-    T_tracer.closeTaskTime("SUCCESS");
-    T_tracer.endTaskTracer();
-
     return new TrainDBListResultSet(header, exportSynopsisInfo);
   }
 

diff --git a/traindb-core/src/main/java/traindb/sql/TrainDBSql.java b/traindb-core/src/main/java/traindb/sql/TrainDBSql.java
@@ -164,7 +164,7 @@ public static TrainDBListResultSet run(TrainDBSqlCommand command, TrainDBSqlRunn
       }
       case EXPORT_SYNOPSIS:
         TrainDBSqlExportSynopsis exportSynopsis = (TrainDBSqlExportSynopsis) command;
-        return runner.exportSynopsis(exportSynopsis.getSynopsisName());
+        return runner.exportSynopsis(exportSynopsis.getSynopsisName(), exportSynopsis.getExportFilename());
       case IMPORT_SYNOPSIS:
         TrainDBSqlImportSynopsis importSynopsis = (TrainDBSqlImportSynopsis) command;
         return runner.importSynopsis(importSynopsis.getSynopsisName(),
@@ -396,8 +396,12 @@ public void exitAlterModel(TrainDBSqlParser.AlterModelContext ctx) {
     @Override
     public void exitExportSynopsis(TrainDBSqlParser.ExportSynopsisContext ctx) {
       String synopsisName = ctx.synopsisName().getText();
+      String exportFilename = null;
+      if (ctx.exportToClause() != null) {
+        exportFilename = ctx.exportToClause().filenameString().getText();
+      }
       LOG.debug("EXPORT SYNOPSIS: name=" + synopsisName);
-      commands.add(new TrainDBSqlExportSynopsis(synopsisName));
+      commands.add(new TrainDBSqlExportSynopsis(synopsisName, exportFilename));
     }
 
     @Override

diff --git a/traindb-core/src/main/java/traindb/sql/TrainDBSqlExportSynopsis.java b/traindb-core/src/main/java/traindb/sql/TrainDBSqlExportSynopsis.java
@@ -16,15 +16,21 @@
 
 class TrainDBSqlExportSynopsis extends TrainDBSqlCommand {
   private final String synopsisName;
+  private final String exportFilename;
 
-  TrainDBSqlExportSynopsis(String synopsisName) {
+  TrainDBSqlExportSynopsis(String synopsisName, String exportFilename) {
     this.synopsisName = synopsisName;
+    this.exportFilename = exportFilename;
   }
 
   String getSynopsisName() {
     return synopsisName;
   }
 
+  String getExportFilename() {
+    return exportFilename;
+  }
+
   @Override
   public Type getType() {
     return Type.EXPORT_SYNOPSIS;

diff --git a/traindb-core/src/main/java/traindb/sql/TrainDBSqlRunner.java b/traindb-core/src/main/java/traindb/sql/TrainDBSqlRunner.java
@@ -80,7 +80,7 @@ void createSynopsis(String synopsisName, String modelName, int limitRows, float
 
   void disableModel(String modelName) throws Exception;
 
-  TrainDBListResultSet exportSynopsis(String synopsisName) throws Exception;
+  TrainDBListResultSet exportSynopsis(String synopsisName, String exportFilename) throws Exception;
 
   TrainDBListResultSet importSynopsis(String synopsisName, String synopsisBinaryString)
       throws Exception;