Make models free of XContent

Signed-off-by: Vamsi Manohar <[email protected]>
vmmusings · Apr 16, 2024 · d6cdf3e · d6cdf3e
1 parent 204c7da
commit d6cdf3e
Show file tree

Hide file tree

Showing 16 changed files with 500 additions and 322 deletions.
diff --git a/spark/src/main/java/org/opensearch/sql/spark/asyncquery/model/AsyncQueryJobMetadata.java b/spark/src/main/java/org/opensearch/sql/spark/asyncquery/model/AsyncQueryJobMetadata.java
@@ -7,19 +7,9 @@
 
 package org.opensearch.sql.spark.asyncquery.model;
 
-import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken;
-import static org.opensearch.sql.spark.execution.session.SessionModel.DATASOURCE_NAME;
-import static org.opensearch.sql.spark.execution.statement.StatementModel.QUERY_ID;
-
 import com.google.gson.Gson;
-import java.io.IOException;
-import java.util.Locale;
 import lombok.Data;
 import lombok.EqualsAndHashCode;
-import lombok.SneakyThrows;
-import org.opensearch.core.common.Strings;
-import org.opensearch.core.xcontent.XContentBuilder;
-import org.opensearch.core.xcontent.XContentParser;
 import org.opensearch.index.seqno.SequenceNumbers;
 import org.opensearch.sql.spark.dispatcher.model.JobType;
 import org.opensearch.sql.spark.execution.statestore.StateModel;
@@ -134,29 +124,6 @@ public String toString() {
     return new Gson().toJson(this);
   }
 
-  /**
-   * Converts JobMetadata to XContentBuilder.
-   *
-   * @return XContentBuilder {@link XContentBuilder}
-   * @throws Exception Exception.
-   */
-  @Override
-  public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
-    builder
-        .startObject()
-        .field(QUERY_ID, queryId.getId())
-        .field("type", TYPE_JOBMETA)
-        .field("jobId", jobId)
-        .field("applicationId", applicationId)
-        .field("resultIndex", resultIndex)
-        .field("sessionId", sessionId)
-        .field(DATASOURCE_NAME, datasourceName)
-        .field(JOB_TYPE, jobType.getText().toLowerCase(Locale.ROOT))
-        .field(INDEX_NAME, indexName)
-        .endObject();
-    return builder;
-  }
-
   /** copy builder. update seqNo and primaryTerm */
   public static AsyncQueryJobMetadata copy(
       AsyncQueryJobMetadata copy, long seqNo, long primaryTerm) {
@@ -173,72 +140,6 @@ public static AsyncQueryJobMetadata copy(
         primaryTerm);
   }
 
-  /**
-   * Convert xcontent parser to JobMetadata.
-   *
-   * @param parser parser.
-   * @return JobMetadata {@link AsyncQueryJobMetadata}
-   * @throws IOException IOException.
-   */
-  @SneakyThrows
-  public static AsyncQueryJobMetadata fromXContent(
-      XContentParser parser, long seqNo, long primaryTerm) {
-    AsyncQueryId queryId = null;
-    String jobId = null;
-    String applicationId = null;
-    String resultIndex = null;
-    String sessionId = null;
-    String datasourceName = null;
-    String jobTypeStr = null;
-    String indexName = null;
-    ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.currentToken(), parser);
-    while (!XContentParser.Token.END_OBJECT.equals(parser.nextToken())) {
-      String fieldName = parser.currentName();
-      parser.nextToken();
-      switch (fieldName) {
-        case QUERY_ID:
-          queryId = new AsyncQueryId(parser.textOrNull());
-          break;
-        case "jobId":
-          jobId = parser.textOrNull();
-          break;
-        case "applicationId":
-          applicationId = parser.textOrNull();
-          break;
-        case "resultIndex":
-          resultIndex = parser.textOrNull();
-          break;
-        case "sessionId":
-          sessionId = parser.textOrNull();
-          break;
-        case DATASOURCE_NAME:
-          datasourceName = parser.textOrNull();
-        case JOB_TYPE:
-          jobTypeStr = parser.textOrNull();
-        case INDEX_NAME:
-          indexName = parser.textOrNull();
-        case "type":
-          break;
-        default:
-          throw new IllegalArgumentException("Unknown field: " + fieldName);
-      }
-    }
-    if (jobId == null || applicationId == null) {
-      throw new IllegalArgumentException("jobId and applicationId are required fields.");
-    }
-    return new AsyncQueryJobMetadata(
-        queryId,
-        applicationId,
-        jobId,
-        resultIndex,
-        sessionId,
-        datasourceName,
-        Strings.isNullOrEmpty(jobTypeStr) ? null : JobType.fromString(jobTypeStr),
-        indexName,
-        seqNo,
-        primaryTerm);
-  }
-
   @Override
   public String getId() {
     return queryId.docId();

diff --git a/spark/src/main/java/org/opensearch/sql/spark/dispatcher/model/IndexDMLResult.java b/spark/src/main/java/org/opensearch/sql/spark/dispatcher/model/IndexDMLResult.java
@@ -5,24 +5,19 @@
 
 package org.opensearch.sql.spark.dispatcher.model;
 
-import static org.opensearch.sql.spark.execution.session.SessionModel.DATASOURCE_NAME;
-
-import com.google.common.collect.ImmutableList;
-import java.io.IOException;
 import lombok.Data;
 import lombok.EqualsAndHashCode;
-import org.opensearch.core.xcontent.XContentBuilder;
 import org.opensearch.index.seqno.SequenceNumbers;
 import org.opensearch.sql.spark.execution.statestore.StateModel;
 
 /** Plugin create Index DML result. */
 @Data
 @EqualsAndHashCode(callSuper = false)
 public class IndexDMLResult extends StateModel {
-  private static final String QUERY_ID = "queryId";
-  private static final String QUERY_RUNTIME = "queryRunTime";
-  private static final String UPDATE_TIME = "updateTime";
-  private static final String DOC_ID_PREFIX = "index";
+  public static final String QUERY_ID = "queryId";
+  public static final String QUERY_RUNTIME = "queryRunTime";
+  public static final String UPDATE_TIME = "updateTime";
+  public static final String DOC_ID_PREFIX = "index";
 
   private final String queryId;
   private final String status;
@@ -55,20 +50,4 @@ public long getSeqNo() {
   public long getPrimaryTerm() {
     return SequenceNumbers.UNASSIGNED_PRIMARY_TERM;
   }
-
-  @Override
-  public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
-    builder
-        .startObject()
-        .field(QUERY_ID, queryId)
-        .field("status", status)
-        .field("error", error)
-        .field(DATASOURCE_NAME, datasourceName)
-        .field(QUERY_RUNTIME, queryRunTime)
-        .field(UPDATE_TIME, updateTime)
-        .field("result", ImmutableList.of())
-        .field("schema", ImmutableList.of())
-        .endObject();
-    return builder;
-  }
 }
diff --git a/spark/src/main/java/org/opensearch/sql/spark/execution/session/SessionModel.java b/spark/src/main/java/org/opensearch/sql/spark/execution/session/SessionModel.java
@@ -8,13 +8,8 @@
 import static org.opensearch.sql.spark.execution.session.SessionState.NOT_STARTED;
 import static org.opensearch.sql.spark.execution.session.SessionType.INTERACTIVE;
 
-import java.io.IOException;
 import lombok.Builder;
 import lombok.Data;
-import lombok.SneakyThrows;
-import org.opensearch.core.xcontent.XContentBuilder;
-import org.opensearch.core.xcontent.XContentParser;
-import org.opensearch.core.xcontent.XContentParserUtils;
 import org.opensearch.index.seqno.SequenceNumbers;
 import org.opensearch.sql.spark.execution.statestore.StateModel;
 
@@ -48,24 +43,6 @@ public class SessionModel extends StateModel {
   private final long seqNo;
   private final long primaryTerm;
 
-  @Override
-  public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
-    builder
-        .startObject()
-        .field(VERSION, version)
-        .field(TYPE, SESSION_DOC_TYPE)
-        .field(SESSION_TYPE, sessionType.getSessionType())
-        .field(SESSION_ID, sessionId.getSessionId())
-        .field(SESSION_STATE, sessionState.getSessionState())
-        .field(DATASOURCE_NAME, datasourceName)
-        .field(APPLICATION_ID, applicationId)
-        .field(JOB_ID, jobId)
-        .field(LAST_UPDATE_TIME, lastUpdateTime)
-        .field(ERROR, error)
-        .endObject();
-    return builder;
-  }
-
   public static SessionModel of(SessionModel copy, long seqNo, long primaryTerm) {
     return builder()
         .version(copy.version)
@@ -99,52 +76,6 @@ public static SessionModel copyWithState(
         .build();
   }
 
-  @SneakyThrows
-  public static SessionModel fromXContent(XContentParser parser, long seqNo, long primaryTerm) {
-    SessionModelBuilder builder = new SessionModelBuilder();
-    XContentParserUtils.ensureExpectedToken(
-        XContentParser.Token.START_OBJECT, parser.currentToken(), parser);
-    while (!XContentParser.Token.END_OBJECT.equals(parser.nextToken())) {
-      String fieldName = parser.currentName();
-      parser.nextToken();
-      switch (fieldName) {
-        case VERSION:
-          builder.version(parser.text());
-          break;
-        case SESSION_TYPE:
-          builder.sessionType(SessionType.fromString(parser.text()));
-          break;
-        case SESSION_ID:
-          builder.sessionId(new SessionId(parser.text()));
-          break;
-        case SESSION_STATE:
-          builder.sessionState(SessionState.fromString(parser.text()));
-          break;
-        case DATASOURCE_NAME:
-          builder.datasourceName(parser.text());
-          break;
-        case ERROR:
-          builder.error(parser.text());
-          break;
-        case APPLICATION_ID:
-          builder.applicationId(parser.text());
-          break;
-        case JOB_ID:
-          builder.jobId(parser.text());
-          break;
-        case LAST_UPDATE_TIME:
-          builder.lastUpdateTime(parser.longValue());
-          break;
-        case TYPE:
-          // do nothing.
-          break;
-      }
-    }
-    builder.seqNo(seqNo);
-    builder.primaryTerm(primaryTerm);
-    return builder.build();
-  }
-
   public static SessionModel initInteractiveSession(
       String applicationId, String jobId, SessionId sid, String datasourceName) {
     return builder()

diff --git a/spark/src/main/java/org/opensearch/sql/spark/execution/statement/StatementModel.java b/spark/src/main/java/org/opensearch/sql/spark/execution/statement/StatementModel.java
@@ -5,18 +5,10 @@
 
 package org.opensearch.sql.spark.execution.statement;
 
-import static org.opensearch.sql.spark.execution.session.SessionModel.APPLICATION_ID;
-import static org.opensearch.sql.spark.execution.session.SessionModel.DATASOURCE_NAME;
-import static org.opensearch.sql.spark.execution.session.SessionModel.JOB_ID;
 import static org.opensearch.sql.spark.execution.statement.StatementState.WAITING;
 
-import java.io.IOException;
 import lombok.Builder;
 import lombok.Data;
-import lombok.SneakyThrows;
-import org.opensearch.core.xcontent.XContentBuilder;
-import org.opensearch.core.xcontent.XContentParser;
-import org.opensearch.core.xcontent.XContentParserUtils;
 import org.opensearch.index.seqno.SequenceNumbers;
 import org.opensearch.sql.spark.execution.session.SessionId;
 import org.opensearch.sql.spark.execution.statestore.StateModel;
@@ -55,27 +47,6 @@ public class StatementModel extends StateModel {
   private final long seqNo;
   private final long primaryTerm;
 
-  @Override
-  public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
-    builder
-        .startObject()
-        .field(VERSION, version)
-        .field(TYPE, STATEMENT_DOC_TYPE)
-        .field(STATEMENT_STATE, statementState.getState())
-        .field(STATEMENT_ID, statementId.getId())
-        .field(SESSION_ID, sessionId.getSessionId())
-        .field(APPLICATION_ID, applicationId)
-        .field(JOB_ID, jobId)
-        .field(LANG, langType.getText())
-        .field(DATASOURCE_NAME, datasourceName)
-        .field(QUERY, query)
-        .field(QUERY_ID, queryId)
-        .field(SUBMIT_TIME, submitTime)
-        .field(ERROR, error)
-        .endObject();
-    return builder;
-  }
-
   public static StatementModel copy(StatementModel copy, long seqNo, long primaryTerm) {
     return builder()
         .version("1.0")
@@ -115,61 +86,6 @@ public static StatementModel copyWithState(
         .build();
   }
 
-  @SneakyThrows
-  public static StatementModel fromXContent(XContentParser parser, long seqNo, long primaryTerm) {
-    StatementModel.StatementModelBuilder builder = StatementModel.builder();
-    XContentParserUtils.ensureExpectedToken(
-        XContentParser.Token.START_OBJECT, parser.currentToken(), parser);
-    while (!XContentParser.Token.END_OBJECT.equals(parser.nextToken())) {
-      String fieldName = parser.currentName();
-      parser.nextToken();
-      switch (fieldName) {
-        case VERSION:
-          builder.version(parser.text());
-          break;
-        case TYPE:
-          // do nothing
-          break;
-        case STATEMENT_STATE:
-          builder.statementState(StatementState.fromString(parser.text()));
-          break;
-        case STATEMENT_ID:
-          builder.statementId(new StatementId(parser.text()));
-          break;
-        case SESSION_ID:
-          builder.sessionId(new SessionId(parser.text()));
-          break;
-        case APPLICATION_ID:
-          builder.applicationId(parser.text());
-          break;
-        case JOB_ID:
-          builder.jobId(parser.text());
-          break;
-        case LANG:
-          builder.langType(LangType.fromString(parser.text()));
-          break;
-        case DATASOURCE_NAME:
-          builder.datasourceName(parser.text());
-          break;
-        case QUERY:
-          builder.query(parser.text());
-          break;
-        case QUERY_ID:
-          builder.queryId(parser.text());
-          break;
-        case SUBMIT_TIME:
-          builder.submitTime(parser.longValue());
-          break;
-        case ERROR:
-          builder.error(parser.text());
-          break;
-      }
-    }
-    builder.seqNo(seqNo);
-    builder.primaryTerm(primaryTerm);
-    return builder.build();
-  }
-
   public static StatementModel submitStatement(
       SessionId sid,
       String applicationId,

diff --git a/spark/src/main/java/org/opensearch/sql/spark/execution/statestore/CopyBuilder.java b/spark/src/main/java/org/opensearch/sql/spark/execution/statestore/CopyBuilder.java
@@ -0,0 +1,5 @@
+package org.opensearch.sql.spark.execution.statestore;
+
+public interface CopyBuilder<T> {
+  T of(T copy, long seqNo, long primaryTerm);
+}
diff --git a/spark/src/main/java/org/opensearch/sql/spark/execution/statestore/FromXContent.java b/spark/src/main/java/org/opensearch/sql/spark/execution/statestore/FromXContent.java
@@ -0,0 +1,7 @@
+package org.opensearch.sql.spark.execution.statestore;
+
+import org.opensearch.core.xcontent.XContentParser;
+
+public interface FromXContent<T extends StateModel> {
+  T fromXContent(XContentParser parser, long seqNo, long primaryTerm);
+}
diff --git a/spark/src/main/java/org/opensearch/sql/spark/execution/statestore/StateCopyBuilder.java b/spark/src/main/java/org/opensearch/sql/spark/execution/statestore/StateCopyBuilder.java
@@ -0,0 +1,5 @@
+package org.opensearch.sql.spark.execution.statestore;
+
+public interface StateCopyBuilder<T extends StateModel, S> {
+  T of(T copy, S state, long seqNo, long primaryTerm);
+}