Refactor: reuse a query engine object

Author: taewhi

traindb-core/src/main/java/traindb/jdbc/Driver.java

@@ -37,15 +37,16 @@
 import org.apache.calcite.jdbc.CalciteConnection;
 import org.apache.calcite.jdbc.CalciteFactory;
 import org.apache.calcite.jdbc.CalcitePrepare;
+import org.apache.calcite.jdbc.CalcitePrepare.Context;
 import org.apache.calcite.jdbc.CalciteSchema;
-import org.apache.calcite.linq4j.function.Function0;
+import org.apache.calcite.linq4j.function.Function1;
 import org.checkerframework.checker.nullness.qual.Nullable;
 import traindb.prepare.TrainDBPrepareImpl;
 import traindb.schema.SchemaManager;
 
 public class Driver extends UnregisteredDriver {
  public static final String CONNECT_STRING_PREFIX = "jdbc:traindb:";
- Function0<CalcitePrepare> DEFAULT_TRAINDB_FACTORY = TrainDBPrepareImpl::new;
+ Function1<Context, CalcitePrepare> DEFAULT_TRAINDB_FACTORY = TrainDBPrepareImpl::new;
 
  static {
  try {
@@ -56,15 +57,15 @@ public class Driver extends UnregisteredDriver {
  }
  }
 
- final Function0<CalcitePrepare> prepareFactory;
+ final Function1<Context, CalcitePrepare> prepareFactory;
 
  @SuppressWarnings("method.invocation.invalid")
  public Driver() {
  super();
  this.prepareFactory = createPrepareFactory();
  }
 
- protected Function0<CalcitePrepare> createPrepareFactory() {
+ protected Function1<Context, CalcitePrepare> createPrepareFactory() {
  return DEFAULT_TRAINDB_FACTORY;
  }
 

traindb-core/src/main/java/traindb/jdbc/TrainDBConnectionImpl.java

@@ -61,7 +61,7 @@
 import org.apache.calcite.linq4j.Ord;
 import org.apache.calcite.linq4j.QueryProvider;
 import org.apache.calcite.linq4j.Queryable;
-import org.apache.calcite.linq4j.function.Function0;
+import org.apache.calcite.linq4j.function.Function1;
 import org.apache.calcite.linq4j.tree.Expression;
 import org.apache.calcite.materialize.Lattice;
 import org.apache.calcite.materialize.MaterializationService;
@@ -110,7 +110,7 @@ public abstract class TrainDBConnectionImpl
  private SchemaManager schemaManager;
  private BasicDataSource dataSource;
 
- final Function0<CalcitePrepare> prepareFactory;
+ final Function1<Context, CalcitePrepare> prepareFactory;
  final CalciteServer server = new CalciteServerImpl();
  CalciteSchema rootSchema;
 
@@ -369,7 +369,7 @@ <T> CalcitePrepare.CalciteSignature<T> parseQuery(
  Context prepareContext, long maxRowCount) {
  CalcitePrepare.Dummy.push(prepareContext);
  try {
- final CalcitePrepare prepare = prepareFactory.apply();
+ final CalcitePrepare prepare = prepareFactory.apply(prepareContext);
  return prepare.prepareSql(prepareContext, query, Object[].class,
  maxRowCount);
  } finally {

traindb-core/src/main/java/traindb/jdbc/TrainDBStatement.java

@@ -67,7 +67,6 @@ public TrainDBConnectionImpl getConnection() {
  public <T> CalcitePrepare.CalciteSignature<T> prepare(
  Queryable<T> queryable) {
  final TrainDBConnectionImpl calciteConnection = getConnection();
- final CalcitePrepare prepare = calciteConnection.prepareFactory.apply();
  final CalciteServerStatement serverStatement;
  try {
  serverStatement = calciteConnection.server.getStatement(handle);
@@ -76,6 +75,7 @@ public <T> CalcitePrepare.CalciteSignature<T> prepare(
  }
  final CalcitePrepare.Context prepareContext =
  serverStatement.createPrepareContext();
+ final CalcitePrepare prepare = calciteConnection.prepareFactory.apply(prepareContext);
  return prepare.prepareQueryable(prepareContext, queryable);
  }
 

traindb-core/src/main/java/traindb/prepare/TrainDBPrepareImpl.java

@@ -144,7 +144,11 @@ public class TrainDBPrepareImpl extends CalcitePrepareImpl {
  * plan. If not, enumerable convention is the default. */
  public final boolean enableBindable = Hook.ENABLE_BINDABLE.get(false);
 
- public TrainDBPrepareImpl() {
+ private TrainDBQueryEngine queryEngine;
+
+ public TrainDBPrepareImpl(Context context) {
+ queryEngine = new TrainDBQueryEngine(
+ (TrainDBConnectionImpl) context.getDataContext().getQueryProvider());
  }
 
  @Override public ParseResult parse(
@@ -620,19 +624,18 @@ <T> CalciteSignature<T> prepare2_(
  String currentUser = conn.getProperties().getProperty("user");
 
  if (commands != null && commands.size() > 0) {
- TrainDBQueryEngine engine = new TrainDBQueryEngine(conn);
  try {
  // INSERT QUERY LOGS
- engine.insertQueryLogs(currentTime, currentUser, query.sql);
+ queryEngine.insertQueryLogs(currentTime, currentUser, query.sql);
 
  return convertResultToSignature(context, query.sql,
- TrainDBSql.run(commands.get(0), engine));
+ TrainDBSql.run(commands.get(0), queryEngine));
  } catch (Exception e) {
  throw new RuntimeException(
  "failed to run statement: " + query + "\nerror msg: " + e.getMessage());
  } finally {
  try {
- engine.insertTask();
+ queryEngine.insertTask();
  } catch ( Exception e) {
  throw new RuntimeException(
  "failed to run statement: " + query + "\nerror msg: " + e.getMessage());
@@ -652,8 +655,7 @@ <T> CalciteSignature<T> prepare2_(
 
  // INSERT QUERY LOGS
  try {
- TrainDBQueryEngine engine = new TrainDBQueryEngine(conn);
- engine.insertQueryLogs(currentTime, currentUser, query.sql);
+ queryEngine.insertQueryLogs(currentTime, currentUser, query.sql);
  } catch (Exception e) {
  throw new RuntimeException(
  "failed to query logging: " + query + "\nerror msg: " + e.getMessage());