Skip to content

Commit 6d1ec3a

Browse files
jogroganjogrogan
committed
MCP extensions
1 parent c4a5a66 commit 6d1ec3a

File tree

4 files changed

+600
-36
lines changed

4 files changed

+600
-36
lines changed

‎hoptimator-cli/src/main/java/sqlline/HoptimatorAppConfig.java

Lines changed: 129 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package sqlline;
22

3+
import com.linkedin.hoptimator.Database;
34
import java.nio.charset.StandardCharsets;
45
import java.sql.SQLException;
56
import java.util.Arrays;
@@ -11,8 +12,30 @@
1112
import java.util.Properties;
1213
import java.util.Scanner;
1314

15+
import org.apache.calcite.jdbc.CalcitePrepare;
16+
import org.apache.calcite.jdbc.CalciteSchema;
1417
import org.apache.calcite.plan.RelOptTable;
1518
import org.apache.calcite.rel.RelRoot;
19+
import org.apache.calcite.rel.type.RelDataType;
20+
import org.apache.calcite.rel.type.RelDataTypeFactory;
21+
import org.apache.calcite.rel.type.RelDataTypeImpl;
22+
import org.apache.calcite.rel.type.RelDataTypeSystem;
23+
import org.apache.calcite.rel.type.RelProtoDataType;
24+
import org.apache.calcite.schema.SchemaPlus;
25+
import org.apache.calcite.schema.Table;
26+
import org.apache.calcite.sql.SqlCall;
27+
import org.apache.calcite.sql.SqlIdentifier;
28+
import org.apache.calcite.sql.SqlKind;
29+
import org.apache.calcite.sql.SqlNode;
30+
import org.apache.calcite.sql.SqlNodeList;
31+
import org.apache.calcite.sql.SqlSelect;
32+
import org.apache.calcite.sql.ddl.SqlCreateMaterializedView;
33+
import org.apache.calcite.sql.dialect.CalciteSqlDialect;
34+
import org.apache.calcite.sql.fun.SqlStdOperatorTable;
35+
import org.apache.calcite.sql.parser.SqlParserPos;
36+
import org.apache.calcite.sql.type.SqlTypeFactoryImpl;
37+
import org.apache.calcite.util.Pair;
38+
import org.apache.calcite.util.Util;
1639
import org.jline.reader.Completer;
1740

1841
import com.linkedin.hoptimator.Pipeline;
@@ -96,13 +119,29 @@ public void execute(String line, DispatchCallback dispatchCallback) {
96119
String sql = split[1];
97120
HoptimatorConnection conn = (HoptimatorConnection) sqlline.getConnection();
98121
try {
99-
RelRoot root = HoptimatorDriver.convert(conn, sql).root;
122+
String querySql = sql;
123+
SqlCreateMaterializedView create = null;
124+
SqlNode sqlNode = HoptimatorDriver.parseQuery(conn, sql);
125+
if (sqlNode.getKind().belongsTo(SqlKind.DDL)) {
126+
if (sqlNode instanceof SqlCreateMaterializedView) {
127+
create = (SqlCreateMaterializedView) sqlNode;
128+
final SqlNode q = renameColumns(create.columnList, create.query);
129+
querySql = q.toSqlString(CalciteSqlDialect.DEFAULT).getSql();
130+
} else {
131+
sqlline.error("Unsupported DDL statement: " + sql);
132+
dispatchCallback.setToFailure();
133+
return;
134+
}
135+
}
136+
137+
RelRoot root = HoptimatorDriver.convert(conn, querySql).root;
100138
Properties connectionProperties = conn.connectionProperties();
101139
RelOptTable table = root.rel.getTable();
102140
if (table != null) {
103141
connectionProperties.setProperty(DeploymentService.PIPELINE_OPTION, String.join(".", table.getQualifiedName()));
104142
}
105143
PipelineRel.Implementor plan = DeploymentService.plan(root, conn.materializations(), connectionProperties);
144+
setSink(conn, plan, create, querySql);
106145
sqlline.output(plan.sql(conn).apply(SqlDialect.ANSI));
107146
} catch (SQLException e) {
108147
sqlline.error(e);
@@ -237,14 +276,32 @@ public void execute(String line, DispatchCallback dispatchCallback) {
237276
}
238277
String sql = split[1];
239278
HoptimatorConnection conn = (HoptimatorConnection) sqlline.getConnection();
240-
RelRoot root = HoptimatorDriver.convert(conn, sql).root;
241279
try {
280+
String querySql = sql;
281+
SqlCreateMaterializedView create = null;
282+
SqlNode sqlNode = HoptimatorDriver.parseQuery(conn, sql);
283+
if (sqlNode.getKind().belongsTo(SqlKind.DDL)) {
284+
if (sqlNode instanceof SqlCreateMaterializedView) {
285+
create = (SqlCreateMaterializedView) sqlNode;
286+
final SqlNode q = renameColumns(create.columnList, create.query);
287+
querySql = q.toSqlString(CalciteSqlDialect.DEFAULT).getSql();
288+
} else {
289+
sqlline.error("Unsupported DDL statement: " + sql);
290+
dispatchCallback.setToFailure();
291+
return;
292+
}
293+
}
294+
295+
RelRoot root = HoptimatorDriver.convert(conn, querySql).root;
242296
Properties connectionProperties = conn.connectionProperties();
243297
RelOptTable table = root.rel.getTable();
244298
if (table != null) {
245299
connectionProperties.setProperty(DeploymentService.PIPELINE_OPTION, String.join(".", table.getQualifiedName()));
246300
}
247-
Pipeline pipeline = DeploymentService.plan(root, conn.materializations(), connectionProperties).pipeline("sink", conn);
301+
PipelineRel.Implementor plan = DeploymentService.plan(root, conn.materializations(), connectionProperties);
302+
setSink(conn, plan, create, querySql);
303+
String viewName = create == null ? "sink" : viewName(create.name);
304+
Pipeline pipeline = plan.pipeline(viewName, conn);
248305
List<String> specs = new ArrayList<>();
249306
for (Source source : pipeline.sources()) {
250307
specs.addAll(DeploymentService.specify(source, conn));
@@ -325,4 +382,73 @@ public boolean echoToFile() {
325382
private static boolean startsWith(String s, String prefix) {
326383
return s.matches("(?i)" + prefix + ".*");
327384
}
385+
386+
private static void setSink(HoptimatorConnection conn, PipelineRel.Implementor plan,
387+
SqlCreateMaterializedView create, String querySql) throws SQLException {
388+
if (create == null) {
389+
return;
390+
}
391+
final Pair<CalciteSchema, String> pair = schema(conn.createPrepareContext(), create.name);
392+
String database = ((Database) pair.left.schema).databaseName();
393+
final List<String> schemaPath = pair.left.path(null);
394+
List<String> sinkPath = new ArrayList<>(schemaPath);
395+
String sinkName = pair.right.split("\\$", 2)[0];
396+
sinkPath.add(sinkName);
397+
398+
RelDataTypeFactory typeFactory = new SqlTypeFactoryImpl(RelDataTypeSystem.DEFAULT);
399+
CalcitePrepare.AnalyzeViewResult analyzed = HoptimatorDriver.analyzeView(conn, querySql);
400+
RelProtoDataType protoType = RelDataTypeImpl.proto(analyzed.rowType);
401+
RelDataType viewRowType = protoType.apply(typeFactory);
402+
403+
final SchemaPlus schemaPlus = pair.left.plus();
404+
Table sink = schemaPlus.getTable(sinkName);
405+
final RelDataType rowType;
406+
if (sink != null) {
407+
// For "partial views", the sink may already exist. Use the existing row type.
408+
rowType = sink.getRowType(typeFactory);
409+
} else {
410+
// For normal views, we create the sink based on the view row type.
411+
rowType = viewRowType;
412+
}
413+
414+
plan.setSink(database, sinkPath, rowType, Collections.emptyMap());
415+
}
416+
417+
private static SqlNode renameColumns(SqlNodeList columnList, SqlNode query) {
418+
if (columnList == null) {
419+
return query;
420+
}
421+
final SqlParserPos p = query.getParserPosition();
422+
final SqlNodeList selectList = SqlNodeList.SINGLETON_STAR;
423+
final SqlCall from = SqlStdOperatorTable.AS.createCall(p,
424+
Arrays.asList(query, new SqlIdentifier("_", p), columnList));
425+
return new SqlSelect(p, null, selectList, from, null, null, null, null, null, null, null, null, null);
426+
}
427+
428+
private static Pair<CalciteSchema, String> schema(CalcitePrepare.Context context, SqlIdentifier id) {
429+
final String name;
430+
final List<String> path;
431+
if (id.isSimple()) {
432+
path = context.getDefaultSchemaPath();
433+
name = id.getSimple();
434+
} else {
435+
path = Util.skipLast(id.names);
436+
name = Util.last(id.names);
437+
}
438+
CalciteSchema schema = context.getRootSchema();
439+
for (String p : path) {
440+
schema = Objects.requireNonNull(schema).getSubSchema(p, true);
441+
}
442+
return Pair.of(schema, name);
443+
}
444+
445+
private static String viewName(SqlIdentifier id) {
446+
final String name;
447+
if (id.isSimple()) {
448+
name = id.getSimple();
449+
} else {
450+
name = Util.last(id.names);
451+
}
452+
return name;
453+
}
328454
}

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

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
import org.apache.calcite.schema.Table;
2727
import org.apache.calcite.schema.impl.AbstractSchema;
2828
import org.apache.calcite.sql.SqlNode;
29+
import org.apache.calcite.sql.parser.SqlParseException;
2930
import org.apache.calcite.sql.parser.SqlParser;
3031
import org.apache.calcite.sql.type.SqlTypeFactoryImpl;
3132
import org.apache.calcite.util.Util;
@@ -52,6 +53,17 @@ public class HoptimatorDriver implements java.sql.Driver {
5253
}
5354
}
5455

56+
public static SqlNode parseQuery(HoptimatorConnection conn, String sql) {
57+
Prepare prepare = new Prepare(conn);
58+
SqlParser parser = SqlParser.create(sql, prepare.parserConfig());
59+
try {
60+
return parser.parseQuery();
61+
} catch (SqlParseException e) {
62+
throw new RuntimeException(
63+
"Failed to parse: " + e.getMessage(), e);
64+
}
65+
}
66+
5567
public static CalcitePrepare.ConvertResult convert(HoptimatorConnection conn, String sql) {
5668
CalcitePrepare.Context context = conn.createPrepareContext();
5769
return new Prepare(conn).convert(context, sql);

‎hoptimator-mcp-server/build.gradle

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,11 @@ dependencies {
1010
implementation project(':hoptimator-demodb')
1111
implementation project(':hoptimator-kafka')
1212
implementation project(':hoptimator-venice')
13+
implementation project(':hoptimator-api')
14+
implementation project(':hoptimator-jdbc')
15+
implementation project(':hoptimator-util')
16+
17+
implementation libs.calcite.core
1318

1419
implementation libs.gson
1520
implementation libs.jackson.databind
@@ -26,5 +31,5 @@ tasks.withType(JavaCompile).configureEach {
2631
options.compilerArgs << '-Xlint:deprecation'
2732
options.compilerArgs << '-Xlint:unchecked'
2833
}
29-
34+
3035

0 commit comments

Comments
 (0)