Java Code Examples for org.apache.calcite.schema.SchemaPlus#unwrap()

protected boolean shouldVisitFiles(String schemaName, SchemaPlus schemaPlus) {
  if (filter == null) {
    return true;
  }

  AbstractSchema schema;
  try {
    schema = schemaPlus.unwrap(AbstractSchema.class);
  } catch (ClassCastException e) {
    return false;
  }

  if (!(schema instanceof WorkspaceSchemaFactory.WorkspaceSchema)) {
    return false;
  }

  WorkspaceSchemaFactory.WorkspaceSchema wsSchema = (WorkspaceSchemaFactory.WorkspaceSchema) schema;

  Map<String, String> recordValues = new HashMap<>();
  recordValues.put(FILES_COL_SCHEMA_NAME, schemaName);
  recordValues.put(FILES_COL_ROOT_SCHEMA_NAME, wsSchema.getSchemaPath().get(0));
  recordValues.put(FILES_COL_WORKSPACE_NAME, wsSchema.getName());

  return filter.evaluate(recordValues) != Result.FALSE;
}
 

/**
 * Visit the tables in the given schema. The
 * @param  schemaPath  the path to the given schema
 * @param  schema  the given schema
 */
public void visitTables(String schemaPath, SchemaPlus schema) {
  final AbstractSchema drillSchema = schema.unwrap(AbstractSchema.class);
  for (Pair<String, ? extends Table> tableNameToTable : drillSchema.getTablesByNames(schema.getTableNames())) {
    final String tableName = tableNameToTable.getKey();
    final Table table = tableNameToTable.getValue();
    final TableType tableType = table.getJdbcTableType();
    // Visit the table, and if requested ...
    if(shouldVisitTable(schemaPath, tableName, tableType) && visitTable(schemaPath, tableName, table)) {
      // ... do for each of the table's fields.
      final RelDataType tableRow = table.getRowType(new JavaTypeFactoryImpl(DRILL_REL_DATATYPE_SYSTEM));
      for (RelDataTypeField field: tableRow.getFieldList()) {
        if (shouldVisitColumn(schemaPath, tableName, field.getName())) {
          visitField(schemaPath, tableName, field);
        }
      }
    }
  }
}
 

@Override
public void visitFiles(String schemaName, SchemaPlus schemaPlus) {
  try {
    AbstractSchema schema = schemaPlus.unwrap(AbstractSchema.class);
    if (schema instanceof WorkspaceSchemaFactory.WorkspaceSchema) {
      WorkspaceSchemaFactory.WorkspaceSchema wsSchema = (WorkspaceSchemaFactory.WorkspaceSchema) schema;
      String defaultLocation = wsSchema.getDefaultLocation();
      FileSystem fs = wsSchema.getFS();
      boolean recursive = optionManager.getBoolean(ExecConstants.LIST_FILES_RECURSIVELY);
      // add URI to the path to ensure that directory objects are skipped (see S3AFileSystem.listStatus method)
      FileSystemUtil.listAllSafe(fs, new Path(fs.getUri().toString(), defaultLocation), recursive).forEach(
          fileStatus -> records.add(new Records.File(schemaName, wsSchema, fileStatus))
      );
    }
  } catch (ClassCastException | UnsupportedOperationException e) {
    // ignore the exception since either this is not a Drill schema or schema does not support files listing
  }
}
 

protected boolean shouldVisitSchema(String schemaName, SchemaPlus schema) {
  try {
    // if the schema path is null or empty (try for root schema)
    if (schemaName == null || schemaName.isEmpty()) {
      return false;
    }

    AbstractSchema drillSchema = schema.unwrap(AbstractSchema.class);
    if (!drillSchema.showInInformationSchema()) {
      return false;
    }

    if (filter == null) {
      return true;
    }

    final Map<String, String> recordValues =
        ImmutableMap.of(
            CATS_COL_CATALOG_NAME, IS_CATALOG_NAME,
            SHRD_COL_TABLE_SCHEMA, schemaName,
            SCHS_COL_SCHEMA_NAME, schemaName);

    // If the filter evaluates to false then we don't need to visit the schema.
    // For other two results (TRUE, INCONCLUSIVE) continue to visit the schema.
    return filter.evaluate(recordValues) != Result.FALSE;
  } catch(ClassCastException e) {
    // ignore and return true as this is not a Drill schema
  }
  return true;
}
 

@Override
public boolean visitSchema(String schemaName, SchemaPlus schema) {
  AbstractSchema as = schema.unwrap(AbstractSchema.class);
  records.add(new Records.Schema(IS_CATALOG_NAME, schemaName, "<owner>",
                                 as.getTypeName(), as.isMutable()));
  return false;
}
 

private static void addSchemasToCloseList(final SchemaPlus tree, final List<AutoCloseable> toClose) {
  for(String subSchemaName : tree.getSubSchemaNames()) {
    addSchemasToCloseList(tree.getSubSchema(subSchemaName), toClose);
  }

  try {
    AbstractSchema drillSchemaImpl =  tree.unwrap(AbstractSchema.class);
    toClose.add(drillSchemaImpl);
  } catch (ClassCastException e) {
    // Ignore as the SchemaPlus is not an implementation of Drill schema.
  }
}
 

@Override public Table create(SchemaPlus schema, String tableName, Map operand,
    RelDataType rowType) {
  final RedisSchema redisSchema = schema.unwrap(RedisSchema.class);
  final RelProtoDataType protoRowType =
      rowType != null ? RelDataTypeImpl.proto(rowType) : null;
  return RedisTable.create(redisSchema, tableName, operand, protoRowType);
}
 

/**
 * Loads schema factory(storage plugin) for specified {@code schemaName}
 * @param schemaName the name of the schema
 * @param caseSensitive whether matching for the schema name is case sensitive
 */
public void loadSchemaFactory(String schemaName, boolean caseSensitive) {
  try {
    SchemaPlus schemaPlus = this.plus();
    StoragePlugin plugin = getSchemaFactories().getPlugin(schemaName);
    if (plugin != null && plugin.getConfig().isEnabled()) {
      plugin.registerSchemas(schemaConfig, schemaPlus);
      return;
    }

    // Could not find the plugin of schemaName. The schemaName could be `dfs.tmp`, a 2nd level schema under 'dfs'
    List<String> paths = SchemaUtilites.getSchemaPathAsList(schemaName);
    if (paths.size() == 2) {
      plugin = getSchemaFactories().getPlugin(paths.get(0));
      if (plugin == null) {
        return;
      }

      // Looking for the SchemaPlus for the top level (e.g. 'dfs') of schemaName (e.g. 'dfs.tmp')
      SchemaPlus firstLevelSchema = schemaPlus.getSubSchema(paths.get(0));
      if (firstLevelSchema == null) {
        // register schema for this storage plugin to 'this'.
        plugin.registerSchemas(schemaConfig, schemaPlus);
        firstLevelSchema = schemaPlus.getSubSchema(paths.get(0));
      }
      // Load second level schemas for this storage plugin
      List<SchemaPlus> secondLevelSchemas = new ArrayList<>();
      for (String secondLevelSchemaName : firstLevelSchema.getSubSchemaNames()) {
        secondLevelSchemas.add(firstLevelSchema.getSubSchema(secondLevelSchemaName));
      }

      for (SchemaPlus schema : secondLevelSchemas) {
        org.apache.drill.exec.store.AbstractSchema drillSchema;
        try {
          drillSchema = schema.unwrap(org.apache.drill.exec.store.AbstractSchema.class);
        } catch (ClassCastException e) {
          throw new RuntimeException(String.format("Schema '%s' is not expected under root schema", schema.getName()));
        }
        SubSchemaWrapper wrapper = new SubSchemaWrapper(drillSchema);
        schemaPlus.add(wrapper.getName(), wrapper);
      }
    }
  } catch(ExecutionSetupException | IOException ex) {
    logger.warn("Failed to load schema for \"" + schemaName + "\"!", ex);
    // We can't proceed further without a schema, throw a runtime exception.
    UserException.Builder exceptBuilder =
        UserException
            .resourceError(ex)
            .message("Failed to load schema for \"" + schemaName + "\"!")
            .addContext(ex.getClass().getName() + ": " + ex.getMessage())
            .addContext(UserExceptionUtils.getUserHint(ex)); //Provide hint if it exists
    throw exceptBuilder.build(logger);
  }
}