Fix compound identifier bug (#91)

Author: ryannedolan

hoptimator-kafka/src/test/resources/kafka-ddl.id

@@ -12,9 +12,11 @@ spec:
   job:
     entryClass: com.linkedin.hoptimator.flink.runner.FlinkRunner
     args:
-    - CREATE TABLE IF NOT EXISTS `existing-topic-2` (`KEY` VARCHAR, `VALUE` BINARY) WITH ('connector'='kafka', 'properties.bootstrap.servers'='one-kafka-bootstrap.kafka.svc.cluster.local:9094', 'scan.startup.mode'='earliest-offset', 'topic'='existing-topic-2', 'value.format'='json')
-    - CREATE TABLE IF NOT EXISTS `existing-topic-1` (`KEY` VARCHAR, `VALUE` BINARY) WITH ('connector'='kafka', 'properties.bootstrap.servers'='one-kafka-bootstrap.kafka.svc.cluster.local:9094', 'scan.startup.mode'='earliest-offset', 'topic'='existing-topic-1', 'value.format'='json')
-    - INSERT INTO `existing-topic-1` (`KEY`, `VALUE`) SELECT * FROM `existing-topic-2`
+    - CREATE DATABASE IF NOT EXISTS `KAFKA` WITH ()
+    - CREATE TABLE IF NOT EXISTS `KAFKA`.`existing-topic-2` (`KEY` VARCHAR, `VALUE` BINARY) WITH ('connector'='kafka', 'properties.bootstrap.servers'='one-kafka-bootstrap.kafka.svc.cluster.local:9094', 'scan.startup.mode'='earliest-offset', 'topic'='existing-topic-2', 'value.format'='json')
+    - CREATE DATABASE IF NOT EXISTS `KAFKA` WITH ()
+    - CREATE TABLE IF NOT EXISTS `KAFKA`.`existing-topic-1` (`KEY` VARCHAR, `VALUE` BINARY) WITH ('connector'='kafka', 'properties.bootstrap.servers'='one-kafka-bootstrap.kafka.svc.cluster.local:9094', 'scan.startup.mode'='earliest-offset', 'topic'='existing-topic-1', 'value.format'='json')
+    - INSERT INTO `KAFKA`.`existing-topic-1` (`KEY`, `VALUE`) SELECT * FROM `KAFKA`.`existing-topic-2`
     jarURI: file:///opt/hoptimator-flink-runner.jar
     parallelism: 1
     upgradeMode: stateless

hoptimator-util/src/main/java/com/linkedin/hoptimator/util/planner/PipelineRel.java

@@ -148,8 +148,9 @@ public Pipeline pipeline() throws SQLException {
     private ScriptImplementor script() throws SQLException {
       ScriptImplementor script = ScriptImplementor.empty();
       for (Map.Entry<Source, RelDataType> source : sources.entrySet()) {
+        script = script.database(source.getKey().schema());
         Map<String, String> configs = ConnectionService.configure(source.getKey(), Source.class);
-        script = script.connector(source.getKey().table(), source.getValue(), configs);
+        script = script.connector(source.getKey().schema(), source.getKey().table(), source.getValue(), configs);
       }
       return script;
     }
@@ -163,8 +164,9 @@ public Function<SqlDialect, String> sql() throws SQLException {
       }
       Sink sink = new Sink(sinkDatabase, sinkPath, sinkOptions);
       Map<String, String> sinkConfigs = ConnectionService.configure(sink, Sink.class);
-      script = script.connector(sink.table(), targetRowType, sinkConfigs);
-      script = script.insert(sink.table(), query, targetFields);
+      script = script.database(sink.schema());
+      script = script.connector(sink.schema(), sink.table(), targetRowType, sinkConfigs);
+      script = script.insert(sink.schema(), sink.table(), query, targetFields);
       RelOptUtil.equal(sink.table(), targetRowType, "pipeline", query.getRowType(), Litmus.THROW);
       return script.seal();
     }

hoptimator-util/src/main/java/com/linkedin/hoptimator/util/planner/ScriptImplementor.java

@@ -93,8 +93,8 @@ default ScriptImplementor query(RelNode relNode) {
   }
 
   /** Append a connector definition, e.g. `CREATE TABLE ... WITH (...)` */
-  default ScriptImplementor connector(String name, RelDataType rowType, Map<String, String> connectorConfig) {
-    return with(new ConnectorImplementor(name, rowType, connectorConfig));
+  default ScriptImplementor connector(String schema, String table, RelDataType rowType, Map<String, String> connectorConfig) {
+    return with(new ConnectorImplementor(schema, table, rowType, connectorConfig));
   }
 
   /** Append a database definition, e.g. `CREATE DATABASE ...` */
@@ -103,8 +103,8 @@ default ScriptImplementor database(String database) {
   }
 
   /** Append an insert statement, e.g. `INSERT INTO ... SELECT ...` */
-  default ScriptImplementor insert(String name, RelNode relNode, ImmutableList<Pair<Integer, String>> targetFields) {
-    return with(new InsertImplementor(name, relNode, targetFields));
+  default ScriptImplementor insert(String schema, String table, RelNode relNode, ImmutableList<Pair<Integer, String>> targetFields) {
+    return with(new InsertImplementor(schema, table, relNode, targetFields));
   }
 
   /** Render the script as DDL/SQL in the default dialect */
@@ -157,15 +157,9 @@ public void implement(SqlWriter w) {
   /** Implements an arbitrary RelNode as a query */
   class QueryImplementor implements ScriptImplementor {
     private final RelNode relNode;
-    private final boolean dropDatabaseName;
 
     public QueryImplementor(RelNode relNode) {
-      this(relNode, true);
-    }
-
-    public QueryImplementor(RelNode relNode, boolean dropDatabaseName) {
       this.relNode = relNode;
-      this.dropDatabaseName = dropDatabaseName;
     }
 
     @Override
@@ -198,32 +192,24 @@ public SqlNode visit(SqlCall call) {
 
     private static class UnflattenMemberAccess extends SqlShuttle {
       private final Set<String> sinkFieldList;
-      private final boolean dropDatabaseName;
 
       UnflattenMemberAccess(QueryImplementor outer) {
         this.sinkFieldList = outer.relNode.getRowType().getFieldList()
             .stream()
             .map(RelDataTypeField::getName)
             .collect(Collectors.toSet());
-        this.dropDatabaseName = outer.dropDatabaseName;
       }
 
       // SqlShuttle gets called for every field in SELECT and every table name in FROM alike
       // For fields in SELECT, we want to unflatten them as `FOO_BAR`, for tables `FOO.BAR`
-      // or just `BAR` if we need to drop the database name (For Flink)
       @Override
       public SqlNode visit(SqlIdentifier id) {
         if (id.names.size() == 1 && sinkFieldList.contains(id.names.get(0))) {
           id.assignNamesFrom(new SqlIdentifier(id.names.get(0).replaceAll("\\$", "_"), SqlParserPos.ZERO));
         } else {
-          SqlIdentifier replacement;
-          if (id.names.size() == 2 && this.dropDatabaseName) {
-            replacement = new SqlIdentifier(id.names.subList(1, id.names.size()), SqlParserPos.ZERO);
-          } else {
-            replacement = new SqlIdentifier(id.names.stream()
-                .flatMap(x -> Stream.of(x.split("\\$")))
-                .collect(Collectors.toList()), SqlParserPos.ZERO);
-          }
+          SqlIdentifier replacement = new SqlIdentifier(id.names.stream()
+              .flatMap(x -> Stream.of(x.split("\\$")))
+              .collect(Collectors.toList()), SqlParserPos.ZERO);
           id.assignNamesFrom(replacement);
         }
         return id;
@@ -239,20 +225,22 @@ public SqlNode visit(SqlIdentifier id) {
    *  - NULL fields are promoted to BYTES
    */
   class ConnectorImplementor implements ScriptImplementor {
-    private final String name;
+    private final String schema;
+    private final String table;
     private final RelDataType rowType;
     private final Map<String, String> connectorConfig;
 
-    public ConnectorImplementor(String name, RelDataType rowType, Map<String, String> connectorConfig) {
-      this.name = name;
+    public ConnectorImplementor(String schema, String table, RelDataType rowType, Map<String, String> connectorConfig) {
+      this.schema = schema;
+      this.table = table;
       this.rowType = rowType;
       this.connectorConfig = connectorConfig;
     }
 
     @Override
     public void implement(SqlWriter w) {
       w.keyword("CREATE TABLE IF NOT EXISTS");
-      (new IdentifierImplementor(name)).implement(w);
+      (new CompoundIdentifierImplementor(schema, table)).implement(w);
       SqlWriter.Frame frame1 = w.startList("(", ")");
       (new RowTypeSpecImplementor(rowType)).implement(w);
       if (rowType.getField("PRIMARY_KEY", true, false) != null) {
@@ -296,20 +284,22 @@ public void implement(SqlWriter w) {
    *
    * */
   class InsertImplementor implements ScriptImplementor {
-    private final String name;
+    private final String schema;
+    private final String table;
     private final RelNode relNode;
     private final ImmutableList<Pair<Integer, String>> targetFields;
 
-    public InsertImplementor(String name, RelNode relNode, ImmutableList<Pair<Integer, String>> targetFields) {
-      this.name = name;
+    public InsertImplementor(String schema, String table, RelNode relNode, ImmutableList<Pair<Integer, String>> targetFields) {
+      this.schema = schema;
+      this.table = table;
       this.relNode = relNode;
       this.targetFields = targetFields;
     }
 
     @Override
     public void implement(SqlWriter w) {
       w.keyword("INSERT INTO");
-      (new IdentifierImplementor(name)).implement(w);
+      (new CompoundIdentifierImplementor(schema, table)).implement(w);
       RelNode project = dropFields(relNode, targetFields);
       (new ColumnListImplementor(project.getRowType())).implement(w);
       (new QueryImplementor(project)).implement(w);
@@ -394,17 +384,17 @@ public void implement(SqlWriter w) {
 
   /** Implements an identifier like TRACKING.'PageViewEvent' */
   class CompoundIdentifierImplementor implements ScriptImplementor {
-    private final String database;
-    private final String name;
+    private final String schema;
+    private final String table;
 
-    public CompoundIdentifierImplementor(String database, String name) {
-      this.database = database;
-      this.name = name;
+    public CompoundIdentifierImplementor(String schema, String table) {
+      this.schema = schema;
+      this.table = table;
     }
 
     @Override
     public void implement(SqlWriter w) {
-      SqlIdentifier identifier = new SqlIdentifier(Arrays.asList(new String[]{database, name}), SqlParserPos.ZERO);
+      SqlIdentifier identifier = new SqlIdentifier(Arrays.asList(new String[]{schema, table}), SqlParserPos.ZERO);
       identifier.unparse(w, 0, 0);
     }
   }

hoptimator-util/src/test/java/com/linkedin/hoptimator/util/TestDataTypeUtils.java

@@ -41,13 +41,13 @@ public void flattenUnflatten() {
         flattenedNames);
     RelDataType unflattenedType = DataTypeUtils.unflatten(flattenedType, typeFactory);
     RelOptUtil.eq("original", rowType, "flattened-unflattened", unflattenedType, Litmus.THROW);
-    String originalConnector = new ScriptImplementor.ConnectorImplementor("T1",
+    String originalConnector = new ScriptImplementor.ConnectorImplementor("S", "T1",
         rowType, Collections.emptyMap()).sql();
-    String unflattenedConnector = new ScriptImplementor.ConnectorImplementor("T1",
+    String unflattenedConnector = new ScriptImplementor.ConnectorImplementor("S", "T1",
         unflattenedType, Collections.emptyMap()).sql();
     Assertions.assertEquals(originalConnector, unflattenedConnector,
         "Flattening and unflattening data types should have no impact on connector");
-    Assertions.assertEquals("CREATE TABLE IF NOT EXISTS `T1` (`FOO` ROW(`QUX` VARCHAR, "
+    Assertions.assertEquals("CREATE TABLE IF NOT EXISTS `S`.`T1` (`FOO` ROW(`QUX` VARCHAR, "
         + "`QIZ` VARCHAR), `BAR` ROW(`BAZ` VARCHAR)) WITH ();", unflattenedConnector,
         "Flattened-unflattened connector should be correct");
   }
@@ -73,9 +73,9 @@ public void flattenNestedArrays() {
         .collect(Collectors.toList());
     Assertions.assertIterableEquals(Arrays.asList(new String[]{"FOO", "BAR", "CAR"}),
         flattenedNames);
-    String flattenedConnector = new ScriptImplementor.ConnectorImplementor("T1",
+    String flattenedConnector = new ScriptImplementor.ConnectorImplementor("S", "T1",
         flattenedType, Collections.emptyMap()).sql();
-    Assertions.assertEquals("CREATE TABLE IF NOT EXISTS `T1` (`FOO` ANY ARRAY, "
+    Assertions.assertEquals("CREATE TABLE IF NOT EXISTS `S`.`T1` (`FOO` ANY ARRAY, "
         + "`BAR` ANY ARRAY, `CAR` FLOAT ARRAY) WITH ();", flattenedConnector,
         "Flattened connector should have simplified arrays");
   }

hoptimator-venice/src/test/resources/venice-ddl-insert-all.id

@@ -12,9 +12,11 @@ spec:
   job:
     entryClass: com.linkedin.hoptimator.flink.runner.FlinkRunner
     args:
-    - CREATE TABLE IF NOT EXISTS `test-store` (`intField` INTEGER, `stringField` VARCHAR, `KEY` ROW(`id` INTEGER)) WITH ('connector'='venice', 'key.fields'='KEY_id', 'key.fields-prefix'='KEY_', 'key.type'='RECORD', 'partial-update-mode'='true', 'storeName'='test-store', 'value.fields-include'='EXCEPT_KEY')
-    - CREATE TABLE IF NOT EXISTS `test-store-1` (`intField` INTEGER, `stringField` VARCHAR, `KEY_id` INTEGER) WITH ('connector'='venice', 'key.fields'='KEY_id', 'key.fields-prefix'='KEY_', 'key.type'='RECORD', 'partial-update-mode'='true', 'storeName'='test-store-1', 'value.fields-include'='EXCEPT_KEY')
-    - INSERT INTO `test-store-1` (`intField`, `stringField`, `KEY_id`) SELECT * FROM `test-store`
+    - CREATE DATABASE IF NOT EXISTS `VENICE-CLUSTER0` WITH ()
+    - CREATE TABLE IF NOT EXISTS `VENICE-CLUSTER0`.`test-store` (`intField` INTEGER, `stringField` VARCHAR, `KEY` ROW(`id` INTEGER)) WITH ('connector'='venice', 'key.fields'='KEY_id', 'key.fields-prefix'='KEY_', 'key.type'='RECORD', 'partial-update-mode'='true', 'storeName'='test-store', 'value.fields-include'='EXCEPT_KEY')
+    - CREATE DATABASE IF NOT EXISTS `VENICE-CLUSTER0` WITH ()
+    - CREATE TABLE IF NOT EXISTS `VENICE-CLUSTER0`.`test-store-1` (`intField` INTEGER, `stringField` VARCHAR, `KEY_id` INTEGER) WITH ('connector'='venice', 'key.fields'='KEY_id', 'key.fields-prefix'='KEY_', 'key.type'='RECORD', 'partial-update-mode'='true', 'storeName'='test-store-1', 'value.fields-include'='EXCEPT_KEY')
+    - INSERT INTO `VENICE-CLUSTER0`.`test-store-1` (`intField`, `stringField`, `KEY_id`) SELECT * FROM `VENICE-CLUSTER0`.`test-store`
     jarURI: file:///opt/hoptimator-flink-runner.jar
     parallelism: 1
     upgradeMode: stateless

hoptimator-venice/src/test/resources/venice-ddl-insert-partial.id

@@ -12,9 +12,11 @@ spec:
   job:
     entryClass: com.linkedin.hoptimator.flink.runner.FlinkRunner
     args:
-    - CREATE TABLE IF NOT EXISTS `test-store` (`intField` INTEGER, `stringField` VARCHAR, `KEY` ROW(`id` INTEGER)) WITH ('connector'='venice', 'key.fields'='KEY_id', 'key.fields-prefix'='KEY_', 'key.type'='RECORD', 'partial-update-mode'='true', 'storeName'='test-store', 'value.fields-include'='EXCEPT_KEY')
-    - CREATE TABLE IF NOT EXISTS `test-store-1` (`intField` INTEGER, `stringField` VARCHAR, `KEY_id` INTEGER) WITH ('connector'='venice', 'key.fields'='KEY_id', 'key.fields-prefix'='KEY_', 'key.type'='RECORD', 'partial-update-mode'='true', 'storeName'='test-store-1', 'value.fields-include'='EXCEPT_KEY')
-    - INSERT INTO `test-store-1` (`intField`, `KEY_id`) SELECT CAST(`stringField` AS SIGNED) AS `intField`, `KEY_id` FROM `test-store`
+    - CREATE DATABASE IF NOT EXISTS `VENICE-CLUSTER0` WITH ()
+    - CREATE TABLE IF NOT EXISTS `VENICE-CLUSTER0`.`test-store` (`intField` INTEGER, `stringField` VARCHAR, `KEY` ROW(`id` INTEGER)) WITH ('connector'='venice', 'key.fields'='KEY_id', 'key.fields-prefix'='KEY_', 'key.type'='RECORD', 'partial-update-mode'='true', 'storeName'='test-store', 'value.fields-include'='EXCEPT_KEY')
+    - CREATE DATABASE IF NOT EXISTS `VENICE-CLUSTER0` WITH ()
+    - CREATE TABLE IF NOT EXISTS `VENICE-CLUSTER0`.`test-store-1` (`intField` INTEGER, `stringField` VARCHAR, `KEY_id` INTEGER) WITH ('connector'='venice', 'key.fields'='KEY_id', 'key.fields-prefix'='KEY_', 'key.type'='RECORD', 'partial-update-mode'='true', 'storeName'='test-store-1', 'value.fields-include'='EXCEPT_KEY')
+    - INSERT INTO `VENICE-CLUSTER0`.`test-store-1` (`intField`, `KEY_id`) SELECT CAST(`stringField` AS SIGNED) AS `intField`, `KEY_id` FROM `VENICE-CLUSTER0`.`test-store`
     jarURI: file:///opt/hoptimator-flink-runner.jar
     parallelism: 1
     upgradeMode: stateless

hoptimator-venice/src/test/resources/venice-ddl-select.id

@@ -12,9 +12,11 @@ spec:
   job:
     entryClass: com.linkedin.hoptimator.flink.runner.FlinkRunner
     args:
-    - CREATE TABLE IF NOT EXISTS `test-store-1` (`intField` INTEGER, `stringField` VARCHAR, `KEY` ROW(`id` INTEGER)) WITH ('connector'='venice', 'key.fields'='KEY_id', 'key.fields-prefix'='KEY_', 'key.type'='RECORD', 'partial-update-mode'='true', 'storeName'='test-store-1', 'value.fields-include'='EXCEPT_KEY')
-    - CREATE TABLE IF NOT EXISTS `SINK` (`intField` INTEGER, `stringField` VARCHAR, `KEY_id` INTEGER) WITH ()
-    - INSERT INTO `SINK` (`intField`, `stringField`, `KEY_id`) SELECT * FROM `test-store-1`
+    - CREATE DATABASE IF NOT EXISTS `VENICE-CLUSTER0` WITH ()
+    - CREATE TABLE IF NOT EXISTS `VENICE-CLUSTER0`.`test-store-1` (`intField` INTEGER, `stringField` VARCHAR, `KEY` ROW(`id` INTEGER)) WITH ('connector'='venice', 'key.fields'='KEY_id', 'key.fields-prefix'='KEY_', 'key.type'='RECORD', 'partial-update-mode'='true', 'storeName'='test-store-1', 'value.fields-include'='EXCEPT_KEY')
+    - CREATE DATABASE IF NOT EXISTS `PIPELINE` WITH ()
+    - CREATE TABLE IF NOT EXISTS `PIPELINE`.`SINK` (`intField` INTEGER, `stringField` VARCHAR, `KEY_id` INTEGER) WITH ()
+    - INSERT INTO `PIPELINE`.`SINK` (`intField`, `stringField`, `KEY_id`) SELECT * FROM `VENICE-CLUSTER0`.`test-store-1`
     jarURI: file:///opt/hoptimator-flink-runner.jar
     parallelism: 1
     upgradeMode: stateless