Skip to content

Commit f929fd6

Browse files
jogroganjogrogan
authored
Lazy driver loading & disables calcite snapshot behavior (#187)
1 parent 84c65e3 commit f929fd6

File tree

24 files changed

+847
-155
lines changed

24 files changed

+847
-155
lines changed

‎hoptimator-catalog/build.gradle‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,9 +6,11 @@ plugins {
66

77
dependencies {
88
implementation project(':hoptimator-avro')
9+
implementation project(':hoptimator-jdbc')
910
implementation project(':hoptimator-util')
1011
implementation libs.avro
1112
implementation libs.calcite.core
13+
compileOnly libs.findbugs
1214
}
1315

1416
publishing {

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

Lines changed: 28 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -2,36 +2,50 @@
22

33
import java.util.Map;
44
import java.util.stream.Collectors;
5+
import javax.annotation.Nullable;
56

67
import org.apache.calcite.schema.Table;
78
import org.apache.calcite.schema.impl.AbstractSchema;
9+
import org.apache.calcite.schema.lookup.Lookup;
10+
import org.apache.calcite.util.LazyReference;
11+
12+
import com.linkedin.hoptimator.jdbc.schema.LazyTableLookup;
813

914

1015
/** Exposes a Database to Apache Calcite. */
1116
public class DatabaseSchema extends AbstractSchema {
1217
private final Database database;
13-
private final Map<String, Table> tableMap;
18+
private final LazyReference<Lookup<Table>> tables = new LazyReference<>();
1419

15-
public DatabaseSchema(Database database, Map<String, Table> tableMap) {
20+
public DatabaseSchema(Database database) {
1621
this.database = database;
17-
this.tableMap = tableMap;
18-
}
19-
20-
public static DatabaseSchema create(Database database) {
21-
try {
22-
Map<String, Table> tableMap = database.tables().stream().collect(Collectors.toMap(x -> x, x -> new ProtoTable(x, database)));
23-
return new DatabaseSchema(database, tableMap);
24-
} catch (Exception e) {
25-
throw new RuntimeException(e);
26-
}
2722
}
2823

2924
public Database database() {
3025
return database;
3126
}
3227

3328
@Override
34-
public Map<String, Table> getTableMap() {
35-
return tableMap;
29+
public Lookup<Table> tables() {
30+
return tables.getOrCompute(() -> new LazyTableLookup<>() {
31+
32+
@Override
33+
protected Map<String, Table> loadAllTables() throws Exception {
34+
return database.tables().stream().collect(Collectors.toMap(x -> x, x -> new ProtoTable(x, database)));
35+
}
36+
37+
@Override
38+
protected @Nullable Table loadTable(String name) throws Exception {
39+
if (database.tables().contains(name)) {
40+
return new ProtoTable(name, database);
41+
}
42+
return null;
43+
}
44+
45+
@Override
46+
protected String getSchemaDescription() {
47+
return "Database: " + database.toString();
48+
}
49+
});
3650
}
3751
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,6 @@ public Schema create(SchemaPlus parentSchema, String name, Map<String, Object> o
4444
.with("fields.NAME.length", "5")
4545
.with("fields.CEO.length", "5")
4646
.config("COMPANY")));
47-
return DatabaseSchema.create(new Database(name, datagenTables));
47+
return new DatabaseSchema(new Database(name, datagenTables));
4848
}
4949
}

‎hoptimator-demodb/build.gradle‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,10 @@ plugins {
55

66
dependencies {
77
implementation project(':hoptimator-api')
8+
implementation project(':hoptimator-jdbc')
89
implementation project(':hoptimator-util')
910
implementation libs.calcite.core
11+
compileOnly libs.findbugs
1012
}
1113

1214
publishing {

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

Lines changed: 37 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -2,23 +2,53 @@
22

33
import java.util.HashMap;
44
import java.util.Map;
5-
5+
import javax.annotation.Nullable;
66
import org.apache.calcite.schema.Table;
77
import org.apache.calcite.schema.impl.AbstractSchema;
8+
import org.apache.calcite.schema.lookup.Lookup;
9+
import org.apache.calcite.util.LazyReference;
10+
11+
import com.linkedin.hoptimator.jdbc.schema.LazyTableLookup;
812

913

1014
public class AdsSchema extends AbstractSchema {
1115

12-
private final Map<String, Table> tableMap = new HashMap<>();
16+
private final LazyReference<Lookup<Table>> tables = new LazyReference<>();
1317

1418
public AdsSchema() {
15-
tableMap.put("PAGE_VIEWS", new PageViewTable());
16-
tableMap.put("AD_CLICKS", new AdClickTable());
17-
tableMap.put("CAMPAIGNS", new CampaignTable());
1819
}
1920

2021
@Override
21-
public Map<String, Table> getTableMap() {
22-
return tableMap;
22+
public Lookup<Table> tables() {
23+
return tables.getOrCompute(() -> new LazyTableLookup<>() {
24+
25+
@Override
26+
protected Map<String, Table> loadAllTables() {
27+
Map<String, Table> tableMap = new HashMap<>();
28+
tableMap.put("PAGE_VIEWS", new PageViewTable());
29+
tableMap.put("AD_CLICKS", new AdClickTable());
30+
tableMap.put("CAMPAIGNS", new CampaignTable());
31+
return tableMap;
32+
}
33+
34+
@Override
35+
protected @Nullable Table loadTable(String name) {
36+
switch (name) {
37+
case "PAGE_VIEWS":
38+
return new PageViewTable();
39+
case "AD_CLICKS":
40+
return new AdClickTable();
41+
case "CAMPAIGNS":
42+
return new CampaignTable();
43+
default:
44+
return null;
45+
}
46+
}
47+
48+
@Override
49+
protected String getSchemaDescription() {
50+
return "Demo Ads Schema";
51+
}
52+
});
2353
}
2454
}

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package com.linkedin.hoptimator.demodb;
22

3+
import com.linkedin.hoptimator.jdbc.CalciteDriver;
34
import java.io.IOException;
45
import java.sql.Connection;
56
import java.sql.SQLException;
@@ -14,12 +15,11 @@
1415
import org.apache.calcite.avatica.ConnectStringParser;
1516
import org.apache.calcite.avatica.DriverVersion;
1617
import org.apache.calcite.jdbc.CalciteConnection;
17-
import org.apache.calcite.jdbc.Driver;
1818
import org.apache.calcite.schema.SchemaPlus;
1919

2020

2121
/** JDBC driver with fake in-memory data. */
22-
public class DemoDriver extends Driver {
22+
public class DemoDriver extends CalciteDriver {
2323

2424
static {
2525
new DemoDriver().register();

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

Lines changed: 34 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2,22 +2,50 @@
22

33
import java.util.HashMap;
44
import java.util.Map;
5-
5+
import javax.annotation.Nullable;
66
import org.apache.calcite.schema.Table;
77
import org.apache.calcite.schema.impl.AbstractSchema;
8+
import org.apache.calcite.schema.lookup.Lookup;
9+
import org.apache.calcite.util.LazyReference;
10+
11+
import com.linkedin.hoptimator.jdbc.schema.LazyTableLookup;
812

913

1014
public class ProfileSchema extends AbstractSchema {
1115

12-
private final Map<String, Table> tableMap = new HashMap<>();
16+
private final LazyReference<Lookup<Table>> tables = new LazyReference<>();
1317

1418
public ProfileSchema() {
15-
tableMap.put("MEMBERS", new MemberTable());
16-
tableMap.put("COMPANIES", new CompanyTable());
1719
}
1820

1921
@Override
20-
public Map<String, Table> getTableMap() {
21-
return tableMap;
22+
public Lookup<Table> tables() {
23+
return tables.getOrCompute(() -> new LazyTableLookup<>() {
24+
25+
@Override
26+
protected Map<String, Table> loadAllTables() {
27+
Map<String, Table> tableMap = new HashMap<>();
28+
tableMap.put("MEMBERS", new MemberTable());
29+
tableMap.put("COMPANIES", new CompanyTable());
30+
return tableMap;
31+
}
32+
33+
@Override
34+
protected @Nullable Table loadTable(String name) {
35+
switch (name) {
36+
case "MEMBERS":
37+
return new MemberTable();
38+
case "COMPANIES":
39+
return new CompanyTable();
40+
default:
41+
return null;
42+
}
43+
}
44+
45+
@Override
46+
protected String getSchemaDescription() {
47+
return "Demo Profile Schema";
48+
}
49+
});
2250
}
2351
}

‎hoptimator-jdbc/build.gradle‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ dependencies {
1313
implementation libs.calcite.server
1414
implementation libs.slf4j.api
1515
implementation libs.commons.dbcp2
16-
compileOnly libs.findbugs
16+
compileOnly libs.findbugs
1717

1818
testFixturesImplementation libs.quidem
1919
testFixturesImplementation libs.calcite.core
@@ -22,6 +22,7 @@ dependencies {
2222
testFixturesImplementation libs.junit.jupiter.api
2323

2424
testRuntimeOnly project(':hoptimator-demodb')
25+
testCompileOnly libs.findbugs
2526
}
2627

2728
java {

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

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,61 @@
1+
package com.linkedin.hoptimator.jdbc;
2+
3+
import java.sql.Connection;
4+
import java.sql.SQLException;
5+
import java.util.Properties;
6+
import java.util.function.Supplier;
7+
import org.apache.calcite.avatica.AvaticaConnection;
8+
import org.apache.calcite.avatica.ConnectStringParser;
9+
import org.apache.calcite.jdbc.CalciteFactory;
10+
import org.apache.calcite.jdbc.CalcitePrepare;
11+
import org.apache.calcite.jdbc.CalciteSchema;
12+
import org.apache.calcite.jdbc.Driver;
13+
import org.checkerframework.checker.nullness.qual.Nullable;
14+
15+
import static java.util.Objects.requireNonNull;
16+
17+
18+
// CalciteDriver is an extension of Driver that extends the connect() method to allow for disabling caching.
19+
// Caching enables snapshots via CachingCalciteSchema which has the side effect of preloading all tables for a driver
20+
public class CalciteDriver extends Driver {
21+
22+
public CalciteDriver() {
23+
this(null);
24+
}
25+
26+
protected CalciteDriver(@Nullable Supplier<CalcitePrepare> prepareFactory) {
27+
super(prepareFactory);
28+
}
29+
30+
@Override
31+
public CalciteDriver withPrepareFactory(Supplier<CalcitePrepare> prepareFactory) {
32+
requireNonNull(prepareFactory, "prepareFactory");
33+
if (this.prepareFactory == prepareFactory) {
34+
return this;
35+
}
36+
return new CalciteDriver(prepareFactory);
37+
}
38+
39+
@Override
40+
public Connection connect(String url, Properties info) throws SQLException {
41+
return connect(url, info, false);
42+
}
43+
44+
public Connection connect(String url, Properties info, boolean cache) throws SQLException {
45+
if (!this.acceptsURL(url)) {
46+
return null;
47+
} else {
48+
String prefix = this.getConnectStringPrefix();
49+
50+
assert url.startsWith(prefix);
51+
52+
String urlSuffix = url.substring(prefix.length());
53+
Properties info2 = ConnectStringParser.parse(urlSuffix, info);
54+
CalciteSchema rootSchema = CalciteSchema.createRootSchema(true, cache);
55+
AvaticaConnection connection = ((CalciteFactory) this.factory)
56+
.newConnection(this, this.factory, url, info2, rootSchema, null);
57+
this.handler.onConnectionInit(connection);
58+
return connection;
59+
}
60+
}
61+
}

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

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,6 @@
1717
import org.apache.calcite.avatica.ConnectStringParser;
1818
import org.apache.calcite.jdbc.CalciteConnection;
1919
import org.apache.calcite.jdbc.CalcitePrepare;
20-
import org.apache.calcite.jdbc.Driver;
2120
import org.apache.calcite.prepare.CalcitePrepareImpl;
2221
import org.apache.calcite.rel.type.RelDataType;
2322
import org.apache.calcite.rel.type.RelDataTypeFactory;
@@ -121,7 +120,7 @@ public Connection connect(String url, Properties props) throws SQLException {
121120
// to return our Prepare. But our Prepare requires a HoptimatorConnection, which
122121
// we cannot construct yet.
123122
ConnectionHolder holder = new ConnectionHolder();
124-
connection = new Driver().withPrepareFactory(() -> new Prepare(holder))
123+
connection = new CalciteDriver().withPrepareFactory(() -> new Prepare(holder))
125124
.connect("jdbc:calcite:", properties);
126125
if (connection == null) {
127126
throw new IOException("Could not connect to " + url + ": Could not create Calcite connection.");

0 commit comments

Comments
 (0)