Java Code Examples for org.apache.calcite.jdbc.CalcitePrepare#Context

/** Parses and validates a SQL query. For use within Calcite only. */
public static CalcitePrepare.ParseResult parse(
    final CalciteConnection connection, final CalciteSchema schema,
    final List<String> schemaPath, final String sql) {
  final CalcitePrepare prepare = CalcitePrepare.DEFAULT_FACTORY.apply();
  final ImmutableMap<CalciteConnectionProperty, String> propValues =
      ImmutableMap.of();
  final CalcitePrepare.Context context =
      makeContext(connection, schema, schemaPath, null, propValues);
  CalcitePrepare.Dummy.push(context);
  try {
    return prepare.parse(context, sql);
  } finally {
    CalcitePrepare.Dummy.pop(context);
  }
}
 

/** Executes a prepare action. */
public <R> R perform(CalciteServerStatement statement,
    FrameworkConfig config, Frameworks.BasePrepareAction<R> action) {
  final CalcitePrepare.Context prepareContext =
      statement.createPrepareContext();
  final JavaTypeFactory typeFactory = prepareContext.getTypeFactory();
  final CalciteSchema schema =
      config.getDefaultSchema() != null
          ? CalciteSchema.from(config.getDefaultSchema())
          : prepareContext.getRootSchema();
  CalciteCatalogReader catalogReader =
      new CalciteCatalogReader(schema.root(),
          schema.path(null),
          typeFactory,
          prepareContext.config());
  final RexBuilder rexBuilder = new RexBuilder(typeFactory);
  final RelOptPlanner planner =
      createPlanner(prepareContext,
          config.getContext(),
          config.getCostFactory());
  final RelOptCluster cluster = createCluster(planner, rexBuilder);
  return action.apply(cluster, catalogReader,
      prepareContext.getRootSchema().plus(), statement);
}
 

/** Executes a {@code CREATE VIEW} command. */
public void execute(SqlCreateView create,
    CalcitePrepare.Context context) {
  final Pair<CalciteSchema, String> pair =
      schema(context, true, create.name);
  final SchemaPlus schemaPlus = pair.left.plus();
  for (Function function : schemaPlus.getFunctions(pair.right)) {
    if (function.getParameters().isEmpty()) {
      if (!create.getReplace()) {
        throw SqlUtil.newContextException(create.name.getParserPosition(),
            RESOURCE.viewExists(pair.right));
      }
      pair.left.removeFunction(pair.right);
    }
  }
  final SqlNode q = renameColumns(create.columnList, create.query);
  final String sql = q.toSqlString(CalciteSqlDialect.DEFAULT).getSql();
  final ViewTableMacro viewTableMacro =
      ViewTable.viewMacro(schemaPlus, sql, pair.left.path(null),
          context.getObjectPath(), false);
  final TranslatableTable x = viewTableMacro.apply(ImmutableList.of());
  Util.discard(x);
  schemaPlus.add(pair.right, viewTableMacro);
}
 

/** Executes a {@code DROP MATERIALIZED VIEW} command. */
public void execute(SqlDropMaterializedView drop,
    CalcitePrepare.Context context) {
  final Pair<CalciteSchema, String> pair = schema(context, true, drop.name);
  final Table table = pair.left.plus().getTable(pair.right);
  if (table != null) {
    // Materialized view exists.
    execute((SqlDropObject) drop, context);
    if (table instanceof Wrapper) {
      final MaterializationKey materializationKey =
          ((Wrapper) table).unwrap(MaterializationKey.class);
      if (materializationKey != null) {
        MaterializationService.instance()
            .removeMaterialization(materializationKey);
      }
    }
  }
}
 

public CalcitePrepare.Context getPrepareContext() {
  return new CalcitePrepare.Context() {

    @Override
    public JavaTypeFactory getTypeFactory() {
      return typeFactory;
    }

    @Override
    public CalciteSchema getRootSchema() {
      return CalciteSchema.from(rootSchema);
    }

    @Override
    public List<String> getDefaultSchemaPath() {
      return defaultSchema;
    }

    @Override
    public CalciteConnectionConfig config() {
      return cfg;
    }

    @Override
    public CalcitePrepare.SparkHandler spark() {
      return CalcitePrepare.Dummy.getSparkHandler(false);
    }

    @Override
    public DataContext getDataContext() {
      return null;
    }
  };
}
 

/** Populates the table called {@code name} by executing {@code query}. */
protected static void populate(SqlIdentifier name, SqlNode query,
    CalcitePrepare.Context context) {
  // Generate, prepare and execute an "INSERT INTO table query" statement.
  // (It's a bit inefficient that we convert from SqlNode to SQL and back
  // again.)
  final FrameworkConfig config = Frameworks.newConfigBuilder()
      .defaultSchema(
          Objects.requireNonNull(
              Schemas.subSchema(context.getRootSchema(),
                  context.getDefaultSchemaPath())).plus())
      .build();
  final Planner planner = Frameworks.getPlanner(config);
  try {
    final StringBuilder buf = new StringBuilder();
    final SqlPrettyWriter w =
        new SqlPrettyWriter(
            SqlPrettyWriter.config()
                .withDialect(CalciteSqlDialect.DEFAULT)
                .withAlwaysUseParentheses(false),
            buf);
    buf.append("INSERT INTO ");
    name.unparse(w, 0, 0);
    buf.append(" ");
    query.unparse(w, 0, 0);
    final String sql = buf.toString();
    final SqlNode query1 = planner.parse(sql);
    final SqlNode query2 = planner.validate(query1);
    final RelRoot r = planner.rel(query2);
    final PreparedStatement prepare = context.getRelRunner().prepare(r.rel);
    int rowCount = prepare.executeUpdate();
    Util.discard(rowCount);
    prepare.close();
  } catch (SqlParseException | ValidationException
      | RelConversionException | SQLException e) {
    throw new RuntimeException(e);
  }
}
 

/** Populates the table called {@code name} by executing {@code query}. */
static void populate(SqlIdentifier name, SqlNode query,
    CalcitePrepare.Context context) {
  // Generate, prepare and execute an "INSERT INTO table query" statement.
  // (It's a bit inefficient that we convert from SqlNode to SQL and back
  // again.)
  final FrameworkConfig config = Frameworks.newConfigBuilder()
      .defaultSchema(context.getRootSchema().plus())
      .build();
  final Planner planner = Frameworks.getPlanner(config);
  try {
    final StringBuilder buf = new StringBuilder();
    final SqlWriterConfig writerConfig =
        SqlPrettyWriter.config().withAlwaysUseParentheses(false);
    final SqlPrettyWriter w = new SqlPrettyWriter(writerConfig, buf);
    buf.append("INSERT INTO ");
    name.unparse(w, 0, 0);
    buf.append(' ');
    query.unparse(w, 0, 0);
    final String sql = buf.toString();
    final SqlNode query1 = planner.parse(sql);
    final SqlNode query2 = planner.validate(query1);
    final RelRoot r = planner.rel(query2);
    final PreparedStatement prepare = context.getRelRunner().prepare(r.rel);
    int rowCount = prepare.executeUpdate();
    Util.discard(rowCount);
    prepare.close();
  } catch (SqlParseException | ValidationException
      | RelConversionException | SQLException e) {
    throw new RuntimeException(e);
  }
}
 

private static RelNode run(PropAction action, RuleSet rules)
    throws Exception {

  FrameworkConfig config = Frameworks.newConfigBuilder()
      .ruleSets(rules).build();

  final Properties info = new Properties();
  final Connection connection = DriverManager
      .getConnection("jdbc:calcite:", info);
  final CalciteServerStatement statement = connection
      .createStatement().unwrap(CalciteServerStatement.class);
  final CalcitePrepare.Context prepareContext =
        statement.createPrepareContext();
  final JavaTypeFactory typeFactory = prepareContext.getTypeFactory();
  CalciteCatalogReader catalogReader =
        new CalciteCatalogReader(prepareContext.getRootSchema(),
            prepareContext.getDefaultSchemaPath(),
            typeFactory,
            prepareContext.config());
  final RexBuilder rexBuilder = new RexBuilder(typeFactory);
  final RelOptPlanner planner = new VolcanoPlanner(config.getCostFactory(),
      config.getContext());

  // set up rules before we generate cluster
  planner.clearRelTraitDefs();
  planner.addRelTraitDef(RelCollationTraitDef.INSTANCE);
  planner.addRelTraitDef(ConventionTraitDef.INSTANCE);

  planner.clear();
  for (RelOptRule r : rules) {
    planner.addRule(r);
  }

  final RelOptCluster cluster = RelOptCluster.create(planner, rexBuilder);
  return action.apply(cluster, catalogReader,
      prepareContext.getRootSchema().plus());
}
 

/** Executes a {@code DROP SCHEMA} command. */
public void execute(SqlDropSchema drop,
    CalcitePrepare.Context context) {
  final List<String> path = context.getDefaultSchemaPath();
  CalciteSchema schema = context.getRootSchema();
  for (String p : path) {
    schema = schema.getSubSchema(p, true);
  }
  final boolean existed = schema.removeSubSchema(drop.name.getSimple());
  if (!existed && !drop.ifExists) {
    throw SqlUtil.newContextException(drop.name.getParserPosition(),
        RESOURCE.schemaNotFound(drop.name.getSimple()));
  }
}
 

@Override
protected RelOptPlanner createPlanner(
    final CalcitePrepare.Context prepareContext,
    org.apache.calcite.plan.Context externalContext,
    RelOptCostFactory costFactory) {
    RelOptPlanner planner = super.createPlanner(prepareContext, externalContext, costFactory);
    planner.removeRule(EnumerableRules.ENUMERABLE_TABLE_MODIFICATION_RULE);
    planner.addRule(EnumerableTableModifyExtensionRule.INSTANCE);
    return planner;
}
 

/**
 * Returns the SqlValidator with the given {@code context} schema
 * and type factory.
 */
static SqlValidator validator(CalcitePrepare.Context context,
    boolean mutable) {
  return new ContextSqlValidator(context, mutable);
}
 

/** Executes a {@code CREATE FUNCTION} command. */
public void execute(SqlCreateFunction create,
    CalcitePrepare.Context context) {
  throw new UnsupportedOperationException("CREATE FUNCTION is not supported");
}
 

/** Executes {@code DROP FUNCTION}, {@code DROP TABLE},
 * {@code DROP MATERIALIZED VIEW}, {@code DROP TYPE},
 * {@code DROP VIEW} commands. */
public void execute(SqlDropObject drop,
    CalcitePrepare.Context context) {
  final List<String> path = context.getDefaultSchemaPath();
  CalciteSchema schema = context.getRootSchema();
  for (String p : path) {
    schema = schema.getSubSchema(p, true);
  }
  final boolean existed;
  switch (drop.getKind()) {
  case DROP_TABLE:
  case DROP_MATERIALIZED_VIEW:
    existed = schema.removeTable(drop.name.getSimple());
    if (!existed && !drop.ifExists) {
      throw SqlUtil.newContextException(drop.name.getParserPosition(),
          RESOURCE.tableNotFound(drop.name.getSimple()));
    }
    break;
  case DROP_VIEW:
    // Not quite right: removes any other functions with the same name
    existed = schema.removeFunction(drop.name.getSimple());
    if (!existed && !drop.ifExists) {
      throw SqlUtil.newContextException(drop.name.getParserPosition(),
          RESOURCE.viewNotFound(drop.name.getSimple()));
    }
    break;
  case DROP_TYPE:
    existed = schema.removeType(drop.name.getSimple());
    if (!existed && !drop.ifExists) {
      throw SqlUtil.newContextException(drop.name.getParserPosition(),
          RESOURCE.typeNotFound(drop.name.getSimple()));
    }
    break;
  case DROP_FUNCTION:
    existed = schema.removeFunction(drop.name.getSimple());
    if (!existed && !drop.ifExists) {
      throw SqlUtil.newContextException(drop.name.getParserPosition(),
          RESOURCE.functionNotFound(drop.name.getSimple()));
    }
    break;
  case OTHER_DDL:
  default:
    throw new AssertionError(drop.getKind());
  }
}
 

private static CalcitePrepare.Context makeContext(
    final CalciteConnectionConfig connectionConfig,
    final JavaTypeFactory typeFactory,
    final DataContext dataContext,
    final CalciteSchema schema,
    final List<String> schemaPath, final List<String> objectPath_) {
  final ImmutableList<String> objectPath =
      objectPath_ == null ? null : ImmutableList.copyOf(objectPath_);
  return new CalcitePrepare.Context() {
    public JavaTypeFactory getTypeFactory() {
      return typeFactory;
    }

    public CalciteSchema getRootSchema() {
      return schema.root();
    }

    public CalciteSchema getMutableRootSchema() {
      return getRootSchema();
    }

    public List<String> getDefaultSchemaPath() {
      // schemaPath is usually null. If specified, it overrides schema
      // as the context within which the SQL is validated.
      if (schemaPath == null) {
        return schema.path(null);
      }
      return schemaPath;
    }

    public List<String> getObjectPath() {
      return objectPath;
    }

    public CalciteConnectionConfig config() {
      return connectionConfig;
    }

    public DataContext getDataContext() {
      return dataContext;
    }

    public RelRunner getRelRunner() {
      throw new UnsupportedOperationException();
    }

    public CalcitePrepare.SparkHandler spark() {
      final boolean enable = config().spark();
      return CalcitePrepare.Dummy.getSparkHandler(enable);
    }
  };
}
 

@Override public void executeDdl(CalcitePrepare.Context context,
    SqlNode node) {
  dispatcher.invoke(node, context);
}
 

private SQLResponse getPrepareOnlySqlResponse(String projectName, String correctedSql, Connection conn,
        Boolean isPushDown, List<List<String>> results, List<SelectedColumnMeta> columnMetas) throws SQLException {

    CalcitePrepareImpl.KYLIN_ONLY_PREPARE.set(true);

    PreparedStatement preparedStatement = null;
    try {
        preparedStatement = conn.prepareStatement(correctedSql);
        throw new IllegalStateException("Should have thrown OnlyPrepareEarlyAbortException");
    } catch (Exception e) {
        Throwable rootCause = ExceptionUtils.getRootCause(e);
        if (rootCause != null && rootCause instanceof OnlyPrepareEarlyAbortException) {
            OnlyPrepareEarlyAbortException abortException = (OnlyPrepareEarlyAbortException) rootCause;
            CalcitePrepare.Context context = abortException.getContext();
            CalcitePrepare.ParseResult preparedResult = abortException.getPreparedResult();
            List<RelDataTypeField> fieldList = preparedResult.rowType.getFieldList();

            CalciteConnectionConfig config = context.config();

            // Fill in selected column meta
            for (int i = 0; i < fieldList.size(); ++i) {

                RelDataTypeField field = fieldList.get(i);
                String columnName = field.getKey();

                if (columnName.startsWith("_KY_")) {
                    continue;
                }
                BasicSqlType basicSqlType = (BasicSqlType) field.getValue();

                columnMetas.add(new SelectedColumnMeta(false, config.caseSensitive(), false, false,
                        basicSqlType.isNullable() ? 1 : 0, true, basicSqlType.getPrecision(), columnName,
                        columnName, null, null, null, basicSqlType.getPrecision(),
                        basicSqlType.getScale() < 0 ? 0 : basicSqlType.getScale(),
                        basicSqlType.getSqlTypeName().getJdbcOrdinal(), basicSqlType.getSqlTypeName().getName(),
                        true, false, false));
            }

        } else {
            throw e;
        }
    } finally {
        CalcitePrepareImpl.KYLIN_ONLY_PREPARE.set(false);
        DBUtils.closeQuietly(preparedStatement);
    }

    return buildSqlResponse(projectName, isPushDown, results, columnMetas);
}
 

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);
  }
}
 

/** Creates a query planner and initializes it with a default set of
 * rules. */
protected RelOptPlanner createPlanner(CalcitePrepare.Context prepareContext) {
  return createPlanner(prepareContext, null, null);
}
 

/** Executes a DDL statement.
 *
 * <p>The statement identified itself as DDL in the
 * {@link org.apache.calcite.jdbc.CalcitePrepare.ParseResult#kind} field. */
void executeDdl(CalcitePrepare.Context context, SqlNode node);
 

/** Template for methods that execute DDL commands.
 *
 * <p>The base implementation throws {@link UnsupportedOperationException}
 * because a {@link SqlNode} is not DDL, but overloaded methods such as
 * {@code public void execute(SqlCreateFoo, CalcitePrepare.Context)} are
 * called via reflection. */
public void execute(SqlNode node, CalcitePrepare.Context context) {
  throw new UnsupportedOperationException("DDL not supported: " + node);
}