Feat: EXPORT MODEL ~ TO FILE

traindb-project · Mar 28, 2024 · 7c05de3 · 7c05de3
1 parent 9f7a6a6
commit 7c05de3
Show file tree

Hide file tree

Showing 5 changed files with 27 additions and 7 deletions.
diff --git a/traindb-core/src/main/antlr4/traindb/sql/TrainDBSql.g4 b/traindb-core/src/main/antlr4/traindb/sql/TrainDBSql.g4
@@ -146,7 +146,7 @@ optionValue
     ;
 
 exportModel
-    : K_EXPORT K_MODEL modelName
+    : K_EXPORT K_MODEL modelName exportToClause?
     ;
 
 importModel

diff --git a/traindb-core/src/main/java/traindb/engine/TrainDBQueryEngine.java b/traindb-core/src/main/java/traindb/engine/TrainDBQueryEngine.java
@@ -1102,7 +1102,8 @@ public void deleteTasks(Integer cnt) throws Exception {
   }
 
   @Override
-  public TrainDBListResultSet exportModel(String modelName) throws Exception {
+  public TrainDBListResultSet exportModel(String modelName, String exportFilename)
+      throws Exception {
     T_tracer.startTaskTracer("export model " + modelName);
 
     T_tracer.openTaskTime("find : model");
@@ -1131,13 +1132,22 @@ public TrainDBListResultSet exportModel(String modelName) throws Exception {
       catalogContext.updateTrainingStatus(modelName, "FINISHED");
     }
 
-    Path outputPath = Paths.get(runner.getModelPath().getParent().toString(), modelName + ".zip");
+    Path outputPath;
+    if (exportFilename != null) {
+      outputPath = Paths.get(exportFilename).toAbsolutePath();
+    } else {
+      outputPath = Paths.get(runner.getModelPath().getParent().toString(), modelName + ".zip");
+    }
     runner.exportModel(outputPath.toString());
 
     ObjectMapper mapper = new ObjectMapper();
     ZipUtils.addNewFileFromStringToZip("export_metadata.json",
         mapper.writeValueAsString(mModel), outputPath);
 
+    if (exportFilename != null) {
+      return TrainDBListResultSet.empty();
+    }
+
     List<String> header = Arrays.asList("export_model");
     List<List<Object>> exportModelInfo = new ArrayList<>();
     ByteArray byteArray = convertFileToByteArray(new File(outputPath.toString()));

diff --git a/traindb-core/src/main/java/traindb/sql/TrainDBSql.java b/traindb-core/src/main/java/traindb/sql/TrainDBSql.java
@@ -144,7 +144,7 @@ public static TrainDBListResultSet run(TrainDBSqlCommand command, TrainDBSqlRunn
         break;
       case EXPORT_MODEL:
         TrainDBSqlExportModel exportModel = (TrainDBSqlExportModel) command;
-        return runner.exportModel(exportModel.getModelName());
+        return runner.exportModel(exportModel.getModelName(), exportModel.getExportFilename());
       case IMPORT_MODEL:
         TrainDBSqlImportModel importModel = (TrainDBSqlImportModel) command;
         return runner.importModel(importModel.getModelName(), importModel.getModelBinaryString());
@@ -375,8 +375,12 @@ public void exitDescribeTable(TrainDBSqlParser.DescribeTableContext ctx) {
     @Override
     public void exitExportModel(TrainDBSqlParser.ExportModelContext ctx) {
       String modelName = ctx.modelName().getText();
+      String exportFilename = null;
+      if (ctx.exportToClause() != null) {
+        exportFilename = ctx.exportToClause().filenameString().getText();
+      }
       LOG.debug("EXPORT MODEL: name=" + modelName);
-      commands.add(new TrainDBSqlExportModel(modelName));
+      commands.add(new TrainDBSqlExportModel(modelName, exportFilename));
     }
 
     @Override

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

diff --git a/traindb-core/src/main/java/traindb/sql/TrainDBSqlRunner.java b/traindb-core/src/main/java/traindb/sql/TrainDBSqlRunner.java
@@ -70,7 +70,7 @@ void createSynopsis(String synopsisName, TrainDBSqlCommand.SynopsisType synopsis
 
   void deleteTasks(Integer cnt) throws Exception;
 
-  TrainDBListResultSet exportModel(String modelName) throws Exception;
+  TrainDBListResultSet exportModel(String modelName, String exportFilename) throws Exception;
 
   TrainDBListResultSet importModel(String modelName, String modelBinaryString) throws Exception;