Java Code Examples for org.apache.calcite.jdbc.CalciteConnection#getRootSchema()

private Connection setup() throws SQLException {
  // Create a jdbc database & table
  final String db = TempDb.INSTANCE.getUrl();
  Connection c1 = DriverManager.getConnection(db, "", "");
  Statement stmt1 = c1.createStatement();
  // This is a table we can join with the emps from the hr schema
  stmt1.execute("create table table1(id integer not null primary key, "
      + "field1 varchar(10))");
  stmt1.execute("insert into table1 values(100, 'foo')");
  stmt1.execute("insert into table1 values(200, 'bar')");
  c1.close();

  // Make a Calcite schema with both a jdbc schema and a non-jdbc schema
  Connection connection = DriverManager.getConnection("jdbc:calcite:");
  CalciteConnection calciteConnection =
      connection.unwrap(CalciteConnection.class);
  SchemaPlus rootSchema = calciteConnection.getRootSchema();
  rootSchema.add("DB",
      JdbcSchema.create(rootSchema, "DB",
          JdbcSchema.dataSource(db, "org.hsqldb.jdbcDriver", "", ""),
          null, null));
  rootSchema.add("hr", new ReflectiveSchema(new JdbcTest.HrSchema()));
  return connection;
}
 

/** Test case for
 * <a href="https://issues.apache.org/jira/browse/CALCITE-2382">[CALCITE-2382]
 * Sub-query lateral joined to table function</a>. */
@Test void testInlineViewLateralTableFunction() throws SQLException {
  try (Connection connection = DriverManager.getConnection("jdbc:calcite:")) {
    CalciteConnection calciteConnection =
        connection.unwrap(CalciteConnection.class);
    SchemaPlus rootSchema = calciteConnection.getRootSchema();
    SchemaPlus schema = rootSchema.add("s", new AbstractSchema());
    final TableFunction table =
        TableFunctionImpl.create(Smalls.GENERATE_STRINGS_METHOD);
    schema.add("GenerateStrings", table);
    Table tbl = new ScannableTableTest.SimpleTable();
    schema.add("t", tbl);

    final String sql = "select *\n"
        + "from (select 5 as f0 from \"s\".\"t\") \"a\",\n"
        + "  lateral table(\"s\".\"GenerateStrings\"(f0)) as t(n, c)\n"
        + "where char_length(c) > 3";
    ResultSet resultSet = connection.createStatement().executeQuery(sql);
    final String expected = "F0=5; N=4; C=abcd\n"
        + "F0=5; N=4; C=abcd\n"
        + "F0=5; N=4; C=abcd\n"
        + "F0=5; N=4; C=abcd\n";
    assertThat(CalciteAssert.toString(resultSet), equalTo(expected));
  }
}
 

/** Creates a config builder that will contain a view, "MYVIEW", and also
 * the SCOTT JDBC schema, whose tables implement
 * {@link org.apache.calcite.schema.TranslatableTable}. */
static Frameworks.ConfigBuilder expandingConfig(Connection connection)
    throws SQLException {
  final CalciteConnection calciteConnection =
      connection.unwrap(CalciteConnection.class);
  final SchemaPlus root = calciteConnection.getRootSchema();
  CalciteAssert.SchemaSpec spec = CalciteAssert.SchemaSpec.SCOTT;
  CalciteAssert.addSchema(root, spec);
  final String viewSql =
      String.format(Locale.ROOT, "select * from \"%s\".\"%s\" where 1=1",
          spec.schemaName, "EMP");

  // create view
  ViewTableMacro macro = ViewTable.viewMacro(root, viewSql,
      Collections.singletonList("test"), Arrays.asList("test", "view"), false);

  // register view (in root schema)
  root.add("MYVIEW", macro);

  return Frameworks.newConfigBuilder().defaultSchema(root);
}
 

/**
 * Tests a table function that implements {@link ScannableTable} and returns
 * a single column.
 */
@Test void testScannableTableFunction()
    throws SQLException, ClassNotFoundException {
  Connection connection = DriverManager.getConnection("jdbc:calcite:");
  CalciteConnection calciteConnection =
      connection.unwrap(CalciteConnection.class);
  SchemaPlus rootSchema = calciteConnection.getRootSchema();
  SchemaPlus schema = rootSchema.add("s", new AbstractSchema());
  final TableFunction table = TableFunctionImpl.create(Smalls.MAZE_METHOD);
  schema.add("Maze", table);
  final String sql = "select *\n"
      + "from table(\"s\".\"Maze\"(5, 3, 1))";
  ResultSet resultSet = connection.createStatement().executeQuery(sql);
  final String result = "S=abcde\n"
      + "S=xyz\n"
      + "S=generate(w=5, h=3, s=1)\n";
  assertThat(CalciteAssert.toString(resultSet), is(result));
}
 

private Connection setupConnectionWithNestedTable() throws SQLException {
  Connection connection =
      DriverManager.getConnection("jdbc:calcite:");
  CalciteConnection calciteConnection =
      connection.unwrap(CalciteConnection.class);
  SchemaPlus rootSchema = calciteConnection.getRootSchema();
  SchemaPlus schema = rootSchema.add("s", new AbstractSchema());
  schema.add("nested", new NestedCollectionTable());
  return connection;
}
 

public Connection apply(Connection connection) throws SQLException {
  if (schema != null) {
    CalciteConnection con = connection.unwrap(CalciteConnection.class);
    SchemaPlus rootSchema = con.getRootSchema();
    rootSchema.add(name, schema);
  }
  connection.setSchema(name);
  return connection;
}
 

@Test void test() throws SQLException, ClassNotFoundException {
  // Create two databases
  // It's two times hsqldb, but imagine they are different rdbms's
  final String db1 = TempDb.INSTANCE.getUrl();
  Connection c1 = DriverManager.getConnection(db1, "", "");
  Statement stmt1 = c1.createStatement();
  stmt1.execute("create table table1(id varchar(10) not null primary key, "
      + "field1 varchar(10))");
  stmt1.execute("insert into table1 values('a', 'aaaa')");
  c1.close();

  final String db2 = TempDb.INSTANCE.getUrl();
  Connection c2 = DriverManager.getConnection(db2, "", "");
  Statement stmt2 = c2.createStatement();
  stmt2.execute("create table table2(id varchar(10) not null primary key, "
      + "field1 varchar(10))");
  stmt2.execute("insert into table2 values('a', 'aaaa')");
  c2.close();

  // Connect via calcite to these databases
  Connection connection = DriverManager.getConnection("jdbc:calcite:");
  CalciteConnection calciteConnection =
      connection.unwrap(CalciteConnection.class);
  SchemaPlus rootSchema = calciteConnection.getRootSchema();
  final DataSource ds1 =
      JdbcSchema.dataSource(db1, "org.hsqldb.jdbcDriver", "", "");
  rootSchema.add("DB1",
      JdbcSchema.create(rootSchema, "DB1", ds1, null, null));
  final DataSource ds2 =
      JdbcSchema.dataSource(db2, "org.hsqldb.jdbcDriver", "", "");
  rootSchema.add("DB2",
      JdbcSchema.create(rootSchema, "DB2", ds2, null, null));

  Statement stmt3 = connection.createStatement();
  ResultSet rs = stmt3.executeQuery("select table1.id, table1.field1 "
      + "from db1.table1 join db2.table2 on table1.id = table2.id");
  assertThat(CalciteAssert.toString(rs), equalTo("ID=a; FIELD1=aaaa\n"));
}
 

public void run() throws ClassNotFoundException, SQLException {
  Class.forName("org.apache.calcite.jdbc.Driver");
  Connection connection =
      DriverManager.getConnection("jdbc:calcite:");
  CalciteConnection calciteConnection =
      connection.unwrap(CalciteConnection.class);
  SchemaPlus rootSchema = calciteConnection.getRootSchema();
  rootSchema.add("hr", new ReflectiveSchema(new Hr()));
  rootSchema.add("foodmart", new ReflectiveSchema(new Foodmart()));
  Statement statement = connection.createStatement();
  ResultSet resultSet =
      statement.executeQuery("select *\n"
          + "from \"foodmart\".\"sales_fact_1997\" as s\n"
          + "join \"hr\".\"emps\" as e\n"
          + "on e.\"empid\" = s.\"cust_id\"");
  final StringBuilder buf = new StringBuilder();
  while (resultSet.next()) {
    int n = resultSet.getMetaData().getColumnCount();
    for (int i = 1; i <= n; i++) {
      buf.append(i > 1 ? "; " : "")
          .append(resultSet.getMetaData().getColumnLabel(i))
          .append("=")
          .append(resultSet.getObject(i));
    }
    System.out.println(buf.toString());
    buf.setLength(0);
  }
  resultSet.close();
  statement.close();
  connection.close();
}
 

@BeforeEach
public void setUp() throws SQLException {
  Connection connection = DriverManager.getConnection("jdbc:calcite:");
  CalciteConnection calciteConnection =
      connection.unwrap(CalciteConnection.class);
  SchemaPlus rootSchema = calciteConnection.getRootSchema();
  rootSchema.add("test", new ReflectiveSchema(new TestSchema()));
  calciteConnection.setSchema("test");
  this.conn = calciteConnection;
}
 

private void adviseSql(int apiVersion, String sql, Consumer<ResultSet> checker)
    throws SQLException {
  Properties info = new Properties();
  if (apiVersion == 1) {
    info.put("lex", "JAVA");
    info.put("quoting", "DOUBLE_QUOTE");
  } else if (apiVersion == 2) {
    info.put("lex", "SQL_SERVER");
    info.put("quoting", "BRACKET");
  }
  Connection connection =
      DriverManager.getConnection("jdbc:calcite:", info);
  CalciteConnection calciteConnection =
      connection.unwrap(CalciteConnection.class);
  SchemaPlus rootSchema = calciteConnection.getRootSchema();
  rootSchema.add("hr", new ReflectiveSchema(new JdbcTest.HrSchema()));
  SchemaPlus schema = rootSchema.add("s", new AbstractSchema());
  calciteConnection.setSchema("hr");
  final TableFunction getHints =
      apiVersion == 1 ? new SqlAdvisorGetHintsFunction() : new SqlAdvisorGetHintsFunction2();
  schema.add("get_hints", getHints);
  String getHintsSql;
  if (apiVersion == 1) {
    getHintsSql = "select id, names, type from table(\"s\".\"get_hints\"(?, ?)) as t";
  } else {
    getHintsSql = "select id, names, type, replacement from table([s].[get_hints](?, ?)) as t";
  }

  PreparedStatement ps = connection.prepareStatement(getHintsSql);
  SqlParserUtil.StringAndPos sap = SqlParserUtil.findPos(sql);
  ps.setString(1, sap.sql);
  ps.setInt(2, sap.cursor);
  final ResultSet resultSet = ps.executeQuery();
  checker.accept(resultSet);
  resultSet.close();
  connection.close();
}
 

/**
 * Creates the post processor routine to be applied against a Connection.
 *
 * <p>Table schema is based on JdbcTest#Employee
 * (refer to {@link JdbcFrontLinqBackTest#mutable}).
 *
 * @param initialData records to be presented in table
 * @return a connection post-processor
 */
private static CalciteAssert.ConnectionPostProcessor makePostProcessor(
    final List<JdbcTest.Employee> initialData) {
  return connection -> {
    CalciteConnection calciteConnection =
        connection.unwrap(CalciteConnection.class);
    SchemaPlus rootSchema = calciteConnection.getRootSchema();
    SchemaPlus mapSchema = rootSchema.add("foo", new AbstractSchema());
    final String tableName = "bar";
    final JdbcTest.AbstractModifiableTable table =
        mutable(tableName, initialData);
    mapSchema.add(tableName, table);
    return calciteConnection;
  };
}
 

public Connection apply(Connection connection) throws SQLException {
  CalciteConnection con = connection.unwrap(CalciteConnection.class);
  SchemaPlus rootSchema = con.getRootSchema();
  switch (schemaSpec) {
  case CLONE_FOODMART:
  case JDBC_FOODMART_WITH_LATTICE:
    addSchema(rootSchema, SchemaSpec.JDBC_FOODMART);
    /* fall through */
  default:
    addSchema(rootSchema, schemaSpec);
  }
  con.setSchema(schemaSpec.schemaName);
  return connection;
}
 

protected ConnectionPostProcessor newSchema(final String schemaName,
    Pair<String, Table>... tables) {
  return connection -> {
    CalciteConnection con = connection.unwrap(CalciteConnection.class);
    SchemaPlus rootSchema = con.getRootSchema();
    SchemaPlus schema = rootSchema.add(schemaName, new AbstractSchema());
    for (Pair<String, Table> t : tables) {
      schema.add(t.left, t.right);
    }
    connection.setSchema(schemaName);
    return connection;
  };
}
 

/** As {@link #testScannableTableFunction()} but with named parameters. */
@Test void testScannableTableFunctionWithNamedParameters()
    throws SQLException, ClassNotFoundException {
  Connection connection = DriverManager.getConnection("jdbc:calcite:");
  CalciteConnection calciteConnection =
      connection.unwrap(CalciteConnection.class);
  SchemaPlus rootSchema = calciteConnection.getRootSchema();
  SchemaPlus schema = rootSchema.add("s", new AbstractSchema());
  final TableFunction table = TableFunctionImpl.create(Smalls.MAZE2_METHOD);
  schema.add("Maze", table);
  final String sql = "select *\n"
      + "from table(\"s\".\"Maze\"(5, 3, 1))";
  final Statement statement = connection.createStatement();
  ResultSet resultSet = statement.executeQuery(sql);
  final String result = "S=abcde\n"
      + "S=xyz\n";
  assertThat(CalciteAssert.toString(resultSet),
      is(result + "S=generate2(w=5, h=3, s=1)\n"));

  final String sql2 = "select *\n"
      + "from table(\"s\".\"Maze\"(WIDTH => 5, HEIGHT => 3, SEED => 1))";
  resultSet = statement.executeQuery(sql2);
  assertThat(CalciteAssert.toString(resultSet),
      is(result + "S=generate2(w=5, h=3, s=1)\n"));

  final String sql3 = "select *\n"
      + "from table(\"s\".\"Maze\"(HEIGHT => 3, WIDTH => 5))";
  resultSet = statement.executeQuery(sql3);
  assertThat(CalciteAssert.toString(resultSet),
      is(result + "S=generate2(w=5, h=3, s=null)\n"));
  connection.close();
}
 

private Connection getConnectionWithMultiplyFunction() throws SQLException {
  Connection connection =
      DriverManager.getConnection("jdbc:calcite:");
  CalciteConnection calciteConnection =
      connection.unwrap(CalciteConnection.class);
  SchemaPlus rootSchema = calciteConnection.getRootSchema();
  SchemaPlus schema = rootSchema.add("s", new AbstractSchema());
  final TableFunction table =
      TableFunctionImpl.create(Smalls.MULTIPLICATION_TABLE_METHOD);
  schema.add("multiplication", table);
  return connection;
}
 

private SchemaPlus createRootSchema(final CalciteConnection calciteConnection) {
    final SchemaPlus rootSchema = calciteConnection.getRootSchema();
    // Add any custom functions here
    return rootSchema;
}
 

/** As {@link #testScannableTableFunction()} but with named parameters. */
@Test void testMultipleScannableTableFunctionWithNamedParameters()
    throws SQLException, ClassNotFoundException {
  try (Connection connection = DriverManager.getConnection("jdbc:calcite:");
       Statement statement = connection.createStatement()) {
    CalciteConnection calciteConnection =
        connection.unwrap(CalciteConnection.class);
    SchemaPlus rootSchema = calciteConnection.getRootSchema();
    SchemaPlus schema = rootSchema.add("s", new AbstractSchema());
    final TableFunction table1 = TableFunctionImpl.create(Smalls.MAZE_METHOD);
    schema.add("Maze", table1);
    final TableFunction table2 = TableFunctionImpl.create(Smalls.MAZE2_METHOD);
    schema.add("Maze", table2);
    final TableFunction table3 = TableFunctionImpl.create(Smalls.MAZE3_METHOD);
    schema.add("Maze", table3);
    final String sql = "select *\n"
        + "from table(\"s\".\"Maze\"(5, 3, 1))";
    ResultSet resultSet = statement.executeQuery(sql);
    final String result = "S=abcde\n"
        + "S=xyz\n";
    assertThat(CalciteAssert.toString(resultSet),
        is(result + "S=generate(w=5, h=3, s=1)\n"));

    final String sql2 = "select *\n"
        + "from table(\"s\".\"Maze\"(WIDTH => 5, HEIGHT => 3, SEED => 1))";
    resultSet = statement.executeQuery(sql2);
    assertThat(CalciteAssert.toString(resultSet),
        is(result + "S=generate2(w=5, h=3, s=1)\n"));

    final String sql3 = "select *\n"
        + "from table(\"s\".\"Maze\"(HEIGHT => 3, WIDTH => 5))";
    resultSet = statement.executeQuery(sql3);
    assertThat(CalciteAssert.toString(resultSet),
        is(result + "S=generate2(w=5, h=3, s=null)\n"));

    final String sql4 = "select *\n"
        + "from table(\"s\".\"Maze\"(FOO => 'a'))";
    resultSet = statement.executeQuery(sql4);
    assertThat(CalciteAssert.toString(resultSet),
        is(result + "S=generate3(foo=a)\n"));
  }
}
 

private void checkInternal(boolean caseSensitive) throws SQLException {
  final SqlNameMatcher nameMatcher =
      SqlNameMatchers.withCaseSensitive(caseSensitive);
  final String schemaName = "MySchema";
  final String funcName = "MyFUNC";
  final String anotherName = "AnotherFunc";

  try (Connection connection = DriverManager.getConnection("jdbc:calcite:")) {
    CalciteConnection calciteConnection =
        connection.unwrap(CalciteConnection.class);
    SchemaPlus rootSchema = calciteConnection.getRootSchema();
    SchemaPlus schema = rootSchema.add(schemaName, new AbstractSchema());
    final TableFunction table = TableFunctionImpl.create(Smalls.MAZE_METHOD);
    schema.add(funcName, table);
    schema.add(anotherName, table);
    final TableFunction table2 =
        TableFunctionImpl.create(Smalls.MAZE3_METHOD);
    schema.add(funcName, table2);

    final CalciteServerStatement statement =
        connection.createStatement().unwrap(CalciteServerStatement.class);
    final CalcitePrepare.Context prepareContext =
        statement.createPrepareContext();
    final JavaTypeFactory typeFactory = prepareContext.getTypeFactory();
    CalciteCatalogReader reader =
        new CalciteCatalogReader(prepareContext.getRootSchema(),
            ImmutableList.of(), typeFactory, prepareContext.config());

    final List<SqlOperator> operatorList = new ArrayList<>();
    SqlIdentifier myFuncIdentifier =
        new SqlIdentifier(Lists.newArrayList(schemaName, funcName), null,
            SqlParserPos.ZERO, null);
    reader.lookupOperatorOverloads(myFuncIdentifier,
        SqlFunctionCategory.USER_DEFINED_TABLE_FUNCTION, SqlSyntax.FUNCTION,
        operatorList, nameMatcher);
    checkFunctionType(2, funcName, operatorList);

    operatorList.clear();
    reader.lookupOperatorOverloads(myFuncIdentifier,
        SqlFunctionCategory.USER_DEFINED_FUNCTION, SqlSyntax.FUNCTION,
        operatorList, nameMatcher);
    checkFunctionType(0, null, operatorList);

    operatorList.clear();
    SqlIdentifier anotherFuncIdentifier =
        new SqlIdentifier(Lists.newArrayList(schemaName, anotherName), null,
            SqlParserPos.ZERO, null);
    reader.lookupOperatorOverloads(anotherFuncIdentifier,
        SqlFunctionCategory.USER_DEFINED_TABLE_FUNCTION, SqlSyntax.FUNCTION,
        operatorList, nameMatcher);
    checkFunctionType(1, anotherName, operatorList);
  }
}
 

@Test void testSchemaConsistency() throws Exception {
  // Create a database
  final String db = TempDb.INSTANCE.getUrl();
  Connection c1 = DriverManager.getConnection(db, "", "");
  Statement stmt1 = c1.createStatement();
  stmt1.execute("create table table1(id varchar(10) not null primary key, "
      + "field1 varchar(10))");

  // Connect via calcite to these databases
  Connection connection = DriverManager.getConnection("jdbc:calcite:");
  CalciteConnection calciteConnection =
      connection.unwrap(CalciteConnection.class);
  SchemaPlus rootSchema = calciteConnection.getRootSchema();
  final DataSource ds =
      JdbcSchema.dataSource(db, "org.hsqldb.jdbcDriver", "", "");
  rootSchema.add("DB", JdbcSchema.create(rootSchema, "DB", ds, null, null));

  Statement stmt3 = connection.createStatement();
  ResultSet rs;

  // fails, table does not exist
  try {
    rs = stmt3.executeQuery("select * from db.table2");
    fail("expected error, got " + rs);
  } catch (SQLException e) {
    assertThat(e.getCause().getCause().getMessage(),
        equalTo("Object 'TABLE2' not found within 'DB'"));
  }

  stmt1.execute("create table table2(id varchar(10) not null primary key, "
      + "field1 varchar(10))");
  stmt1.execute("insert into table2 values('a', 'aaaa')");

  PreparedStatement stmt2 =
      connection.prepareStatement("select * from db.table2");

  stmt1.execute("alter table table2 add column field2 varchar(10)");

  // "field2" not visible to stmt2
  rs = stmt2.executeQuery();
  assertThat(CalciteAssert.toString(rs), equalTo("ID=a; FIELD1=aaaa\n"));

  // "field2" visible to a new query
  rs = stmt3.executeQuery("select * from db.table2");
  assertThat(CalciteAssert.toString(rs),
      equalTo("ID=a; FIELD1=aaaa; FIELD2=null\n"));
  c1.close();
}
 

public RelRoot plan(String query) {
  try {
    Connection connection = DriverManager.getConnection("jdbc:calcite:");
    CalciteConnection calciteConnection = connection.unwrap(CalciteConnection.class);
    SchemaPlus rootSchema = calciteConnection.getRootSchema();
    registerSourceSchemas(rootSchema);

    List<SamzaSqlScalarFunctionImpl> samzaSqlFunctions = udfMetadata.stream()
        .map(x -> new SamzaSqlScalarFunctionImpl(x))
        .collect(Collectors.toList());

    final List<RelTraitDef> traitDefs = new ArrayList<>();

    traitDefs.add(ConventionTraitDef.INSTANCE);
    traitDefs.add(RelCollationTraitDef.INSTANCE);

    List<SqlOperatorTable> sqlOperatorTables = new ArrayList<>();
    sqlOperatorTables.add(new SamzaSqlOperatorTable());
    sqlOperatorTables.add(new SamzaSqlUdfOperatorTable(samzaSqlFunctions));

    // Using lenient so that !=,%,- are allowed.
    FrameworkConfig frameworkConfig = Frameworks.newConfigBuilder()
        .parserConfig(SqlParser.configBuilder()
            .setLex(Lex.JAVA)
            .setConformance(SqlConformanceEnum.LENIENT)
            .setCaseSensitive(false) // Make Udfs case insensitive
            .build())
        .defaultSchema(rootSchema)
        .operatorTable(new ChainedSqlOperatorTable(sqlOperatorTables))
        .sqlToRelConverterConfig(SqlToRelConverter.Config.DEFAULT)
        .traitDefs(traitDefs)
        .context(Contexts.EMPTY_CONTEXT)
        .costFactory(null)
        .build();
    Planner planner = Frameworks.getPlanner(frameworkConfig);

    SqlNode sql = planner.parse(query);
    SqlNode validatedSql = planner.validate(sql);
    RelRoot relRoot = planner.rel(validatedSql);
    LOG.info("query plan:\n" + RelOptUtil.toString(relRoot.rel, SqlExplainLevel.ALL_ATTRIBUTES));
    return relRoot;
  } catch (Exception e) {
    String errorMsg = SamzaSqlValidator.formatErrorString(query, e);
    LOG.error(errorMsg, e);
    throw new SamzaException(errorMsg, e);
  }
}