Lazy driver loading & disables calcite snapshot behavior (#187)

Author: jogrogan

hoptimator-catalog/build.gradle

@@ -6,9 +6,11 @@ plugins {
 
 dependencies {
   implementation project(':hoptimator-avro')
+  implementation project(':hoptimator-jdbc')
   implementation project(':hoptimator-util')
   implementation libs.avro
   implementation libs.calcite.core
+  compileOnly libs.findbugs
 }
 
 publishing {

hoptimator-catalog/src/main/java/com/linkedin/hoptimator/catalog/DatabaseSchema.java

@@ -2,36 +2,50 @@
 
 import java.util.Map;
 import java.util.stream.Collectors;
+import javax.annotation.Nullable;
 
 import org.apache.calcite.schema.Table;
 import org.apache.calcite.schema.impl.AbstractSchema;
+import org.apache.calcite.schema.lookup.Lookup;
+import org.apache.calcite.util.LazyReference;
+
+import com.linkedin.hoptimator.jdbc.schema.LazyTableLookup;
 
 
 /** Exposes a Database to Apache Calcite. */
 public class DatabaseSchema extends AbstractSchema {
   private final Database database;
-  private final Map<String, Table> tableMap;
+  private final LazyReference<Lookup<Table>> tables = new LazyReference<>();
 
-  public DatabaseSchema(Database database, Map<String, Table> tableMap) {
+  public DatabaseSchema(Database database) {
     this.database = database;
-    this.tableMap = tableMap;
-  }
-
-  public static DatabaseSchema create(Database database) {
-    try {
-      Map<String, Table> tableMap = database.tables().stream().collect(Collectors.toMap(x -> x, x -> new ProtoTable(x, database)));
-      return new DatabaseSchema(database, tableMap);
-    } catch (Exception e) {
-      throw new RuntimeException(e);
-    }
   }
 
   public Database database() {
     return database;
   }
 
   @Override
-  public Map<String, Table> getTableMap() {
-    return tableMap;
+  public Lookup<Table> tables() {
+    return tables.getOrCompute(() -> new LazyTableLookup<>() {
+
+      @Override
+      protected Map<String, Table> loadAllTables() throws Exception {
+        return database.tables().stream().collect(Collectors.toMap(x -> x, x -> new ProtoTable(x, database)));
+      }
+
+      @Override
+      protected @Nullable Table loadTable(String name) throws Exception {
+        if (database.tables().contains(name)) {
+          return new ProtoTable(name, database);
+        }
+        return null;
+      }
+
+      @Override
+      protected String getSchemaDescription() {
+        return "Database: " + database.toString();
+      }
+    });
   }
 }

hoptimator-catalog/src/main/java/com/linkedin/hoptimator/catalog/builtin/DatagenSchemaFactory.java

@@ -44,6 +44,6 @@ public Schema create(SchemaPlus parentSchema, String name, Map<String, Object> o
         .with("fields.NAME.length", "5")
         .with("fields.CEO.length", "5")
         .config("COMPANY")));
-    return DatabaseSchema.create(new Database(name, datagenTables));
+    return new DatabaseSchema(new Database(name, datagenTables));
   }
 }

hoptimator-demodb/build.gradle

@@ -5,8 +5,10 @@ plugins {
 
 dependencies {
   implementation project(':hoptimator-api')
+  implementation project(':hoptimator-jdbc')
   implementation project(':hoptimator-util')
   implementation libs.calcite.core
+  compileOnly libs.findbugs
 }
 
 publishing {

hoptimator-demodb/src/main/java/com/linkedin/hoptimator/demodb/AdsSchema.java

@@ -2,23 +2,53 @@
 
 import java.util.HashMap;
 import java.util.Map;
-
+import javax.annotation.Nullable;
 import org.apache.calcite.schema.Table;
 import org.apache.calcite.schema.impl.AbstractSchema;
+import org.apache.calcite.schema.lookup.Lookup;
+import org.apache.calcite.util.LazyReference;
+
+import com.linkedin.hoptimator.jdbc.schema.LazyTableLookup;
 
 
 public class AdsSchema extends AbstractSchema {
 
-  private final Map<String, Table> tableMap = new HashMap<>();
+  private final LazyReference<Lookup<Table>> tables = new LazyReference<>();
 
   public AdsSchema() {
-    tableMap.put("PAGE_VIEWS", new PageViewTable());
-    tableMap.put("AD_CLICKS", new AdClickTable());
-    tableMap.put("CAMPAIGNS", new CampaignTable());
   }
 
   @Override
-  public Map<String, Table> getTableMap() {
-    return tableMap;
+  public Lookup<Table> tables() {
+    return tables.getOrCompute(() -> new LazyTableLookup<>() {
+
+      @Override
+      protected Map<String, Table> loadAllTables() {
+        Map<String, Table> tableMap = new HashMap<>();
+        tableMap.put("PAGE_VIEWS", new PageViewTable());
+        tableMap.put("AD_CLICKS", new AdClickTable());
+        tableMap.put("CAMPAIGNS", new CampaignTable());
+        return tableMap;
+      }
+
+      @Override
+      protected @Nullable Table loadTable(String name) {
+        switch (name) {
+          case "PAGE_VIEWS":
+            return new PageViewTable();
+          case "AD_CLICKS":
+            return new AdClickTable();
+          case "CAMPAIGNS":
+            return new CampaignTable();
+          default:
+            return null;
+        }
+      }
+
+      @Override
+      protected String getSchemaDescription() {
+        return "Demo Ads Schema";
+      }
+    });
   }
 }

hoptimator-demodb/src/main/java/com/linkedin/hoptimator/demodb/DemoDriver.java

@@ -1,5 +1,6 @@
 package com.linkedin.hoptimator.demodb;
 
+import com.linkedin.hoptimator.jdbc.CalciteDriver;
 import java.io.IOException;
 import java.sql.Connection;
 import java.sql.SQLException;
@@ -14,12 +15,11 @@
 import org.apache.calcite.avatica.ConnectStringParser;
 import org.apache.calcite.avatica.DriverVersion;
 import org.apache.calcite.jdbc.CalciteConnection;
-import org.apache.calcite.jdbc.Driver;
 import org.apache.calcite.schema.SchemaPlus;
 
 
 /** JDBC driver with fake in-memory data. */
-public class DemoDriver extends Driver {
+public class DemoDriver extends CalciteDriver {
 
   static {
     new DemoDriver().register();

hoptimator-demodb/src/main/java/com/linkedin/hoptimator/demodb/ProfileSchema.java

@@ -2,22 +2,50 @@
 
 import java.util.HashMap;
 import java.util.Map;
-
+import javax.annotation.Nullable;
 import org.apache.calcite.schema.Table;
 import org.apache.calcite.schema.impl.AbstractSchema;
+import org.apache.calcite.schema.lookup.Lookup;
+import org.apache.calcite.util.LazyReference;
+
+import com.linkedin.hoptimator.jdbc.schema.LazyTableLookup;
 
 
 public class ProfileSchema extends AbstractSchema {
 
-  private final Map<String, Table> tableMap = new HashMap<>();
+  private final LazyReference<Lookup<Table>> tables = new LazyReference<>();
 
   public ProfileSchema() {
-    tableMap.put("MEMBERS", new MemberTable());
-    tableMap.put("COMPANIES", new CompanyTable());
   }
 
   @Override
-  public Map<String, Table> getTableMap() {
-    return tableMap;
+  public Lookup<Table> tables() {
+    return tables.getOrCompute(() -> new LazyTableLookup<>() {
+
+      @Override
+      protected Map<String, Table> loadAllTables() {
+        Map<String, Table> tableMap = new HashMap<>();
+        tableMap.put("MEMBERS", new MemberTable());
+        tableMap.put("COMPANIES", new CompanyTable());
+        return tableMap;
+      }
+
+      @Override
+      protected @Nullable Table loadTable(String name) {
+        switch (name) {
+          case "MEMBERS":
+            return new MemberTable();
+          case "COMPANIES":
+            return new CompanyTable();
+          default:
+            return null;
+        }
+      }
+
+      @Override
+      protected String getSchemaDescription() {
+        return "Demo Profile Schema";
+      }
+    });
   }
 }

hoptimator-jdbc/build.gradle

@@ -13,7 +13,7 @@ dependencies {
   implementation libs.calcite.server
   implementation libs.slf4j.api
   implementation libs.commons.dbcp2
-  compileOnly libs.findbugs 
+  compileOnly libs.findbugs
 
   testFixturesImplementation libs.quidem
   testFixturesImplementation libs.calcite.core
@@ -22,6 +22,7 @@ dependencies {
   testFixturesImplementation libs.junit.jupiter.api
 
   testRuntimeOnly project(':hoptimator-demodb')
+  testCompileOnly libs.findbugs
 }
 
 java {

hoptimator-jdbc/src/main/java/com/linkedin/hoptimator/jdbc/CalciteDriver.java

@@ -0,0 +1,61 @@
+package com.linkedin.hoptimator.jdbc;
+
+import java.sql.Connection;
+import java.sql.SQLException;
+import java.util.Properties;
+import java.util.function.Supplier;
+import org.apache.calcite.avatica.AvaticaConnection;
+import org.apache.calcite.avatica.ConnectStringParser;
+import org.apache.calcite.jdbc.CalciteFactory;
+import org.apache.calcite.jdbc.CalcitePrepare;
+import org.apache.calcite.jdbc.CalciteSchema;
+import org.apache.calcite.jdbc.Driver;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+import static java.util.Objects.requireNonNull;
+
+
+// CalciteDriver is an extension of Driver that extends the connect() method to allow for disabling caching.
+// Caching enables snapshots via CachingCalciteSchema which has the side effect of preloading all tables for a driver
+public class CalciteDriver extends Driver {
+
+  public CalciteDriver() {
+    this(null);
+  }
+
+  protected CalciteDriver(@Nullable Supplier<CalcitePrepare> prepareFactory) {
+    super(prepareFactory);
+  }
+
+  @Override
+  public CalciteDriver withPrepareFactory(Supplier<CalcitePrepare> prepareFactory) {
+    requireNonNull(prepareFactory, "prepareFactory");
+    if (this.prepareFactory == prepareFactory) {
+      return this;
+    }
+    return new CalciteDriver(prepareFactory);
+  }
+
+  @Override
+  public Connection connect(String url, Properties info) throws SQLException {
+    return connect(url, info, false);
+  }
+
+  public Connection connect(String url, Properties info, boolean cache) throws SQLException {
+    if (!this.acceptsURL(url)) {
+      return null;
+    } else {
+      String prefix = this.getConnectStringPrefix();
+
+      assert url.startsWith(prefix);
+
+      String urlSuffix = url.substring(prefix.length());
+      Properties info2 = ConnectStringParser.parse(urlSuffix, info);
+      CalciteSchema rootSchema = CalciteSchema.createRootSchema(true, cache);
+      AvaticaConnection connection = ((CalciteFactory) this.factory)
+          .newConnection(this, this.factory, url, info2, rootSchema, null);
+      this.handler.onConnectionInit(connection);
+      return connection;
+    }
+  }
+}

hoptimator-jdbc/src/main/java/com/linkedin/hoptimator/jdbc/HoptimatorDriver.java

@@ -17,7 +17,6 @@
 import org.apache.calcite.avatica.ConnectStringParser;
 import org.apache.calcite.jdbc.CalciteConnection;
 import org.apache.calcite.jdbc.CalcitePrepare;
-import org.apache.calcite.jdbc.Driver;
 import org.apache.calcite.prepare.CalcitePrepareImpl;
 import org.apache.calcite.rel.type.RelDataType;
 import org.apache.calcite.rel.type.RelDataTypeFactory;
@@ -121,7 +120,7 @@ public Connection connect(String url, Properties props) throws SQLException {
       // to return our Prepare. But our Prepare requires a HoptimatorConnection, which
       // we cannot construct yet.
       ConnectionHolder holder = new ConnectionHolder();
-      connection = new Driver().withPrepareFactory(() -> new Prepare(holder))
+      connection = new CalciteDriver().withPrepareFactory(() -> new Prepare(holder))
           .connect("jdbc:calcite:", properties);
       if (connection == null) {
         throw new IOException("Could not connect to " + url + ": Could not create Calcite connection.");