Java Code Examples for org.apache.avro.SchemaBuilder#FieldBuilder

public static org.apache.avro.Schema toAvroSchema(String tableName, RelDef relDef) {
    SchemaBuilder.FieldAssembler<org.apache.avro.Schema> schemaBuilder =
        SchemaBuilder.record(tableName).fields();
    Map<String, Integer> keyIndices = Ord.zip(relDef.getKeyFields()).stream()
        .collect(Collectors.toMap(o -> o.e, o -> o.i));
    for (Pair<RelDataTypeField, ColumnStrategy> entry
        : Pair.zip(relDef.getRowType().getFieldList(), relDef.getStrategies())) {
        RelDataTypeField field = entry.left;
        ColumnStrategy strategy = entry.right;
        RelDataType type = field.getType();
        SchemaBuilder.FieldBuilder<org.apache.avro.Schema> fieldBuilder = schemaBuilder.name(field.getName());
        Integer keyIndex = keyIndices.get(field.getName());
        if (keyIndex != null) {
            fieldBuilder = fieldBuilder.prop(SQL_KEY_INDEX_PROP, keyIndex);
            // Keys are always NOT NULL
            strategy = ColumnStrategy.NOT_NULL_STRATEGY;
        }
        schemaBuilder = toAvroType(fieldBuilder, type, strategy);
    }
    return schemaBuilder.endRecord();
}
 

private static SchemaBuilder.FieldAssembler<Schema> createAvroFields(
    final ResultSetMetaData meta,
    final SchemaBuilder.FieldAssembler<Schema> builder,
    final boolean useLogicalTypes)
    throws SQLException {

  for (int i = 1; i <= meta.getColumnCount(); i++) {

    String columnName;
    if (meta.getColumnName(i).isEmpty()) {
      columnName = meta.getColumnLabel(i);
    } else {
      columnName = meta.getColumnName(i);
    }

    int columnType = meta.getColumnType(i);
    String typeName = JDBCType.valueOf(columnType).getName();
    SchemaBuilder.FieldBuilder<Schema> field =
        builder
            .name(normalizeForAvro(columnName))
            .doc(String.format("From sqlType %d %s", columnType, typeName))
            .prop("columnName", columnName)
            .prop("sqlCode", String.valueOf(columnType))
            .prop("typeName", typeName);
    fieldAvroType(columnType, meta.getPrecision(i), field, useLogicalTypes);
  }
  return builder;
}
 

@Override
public Schema guessSchema(String soqlQuery) throws IOException {
    SoqlQuery query = SoqlQuery.getInstance();
    query.init(soqlQuery);

    SchemaBuilder.FieldAssembler fieldAssembler = SchemaBuilder.record("GuessedSchema").fields();
    DescribeSObjectResult describeSObjectResult = null;

    try {
        describeSObjectResult = connect(null).connection.describeSObject(query.getDrivingEntityName());
    } catch (ConnectionException e) {
        throw new RuntimeException(e.getMessage());
    }

    Schema entitySchema = SalesforceAvroRegistry.get().inferSchema(describeSObjectResult);

    for (FieldDescription fieldDescription : query.getFieldDescriptions()) {
        Schema.Field schemaField = entitySchema.getField(fieldDescription.getSimpleName());

        SchemaBuilder.FieldBuilder builder = fieldAssembler.name(fieldDescription.getFullName());

        Schema fieldType = null;
        if (schemaField != null) {
            Map<String, Object> props = schemaField.getObjectProps();
            for (Map.Entry<String, Object> entry : props.entrySet()) {
                builder.prop(entry.getKey(), String.valueOf(entry.getValue()));
            }
            fieldType = schemaField.schema();
        } else {
            fieldType = DEFAULT_GUESS_SCHEMA_TYPE;
        }
        builder.type(fieldType).noDefault();
    }

    return (Schema) fieldAssembler.endRecord();
}
 

public Collection<Schema> convert(Ddl ddl) {
  Collection<Schema> schemas = new ArrayList<>();

  for (Table table : ddl.allTables()) {
    SchemaBuilder.RecordBuilder<Schema> recordBuilder =
        SchemaBuilder.record(table.name()).namespace(this.namespace);
    recordBuilder.prop("googleFormatVersion", version);
    recordBuilder.prop("googleStorage", "CloudSpanner");
    if (table.interleaveInParent() != null) {
      recordBuilder.prop("spannerParent", table.interleaveInParent());
      recordBuilder.prop(
          "spannerOnDeleteAction", table.onDeleteCascade() ? "cascade" : "no action");
    }
    if (table.primaryKeys() != null) {
      String encodedPk =
          table.primaryKeys().stream()
              .map(IndexColumn::prettyPrint)
              .collect(Collectors.joining(","));
      recordBuilder.prop("spannerPrimaryKey", encodedPk);
    }
    for (int i = 0; i < table.primaryKeys().size(); i++) {
      recordBuilder.prop("spannerPrimaryKey_" + i, table.primaryKeys().get(i).prettyPrint());
    }
    for (int i = 0; i < table.indexes().size(); i++) {
      recordBuilder.prop("spannerIndex_" + i, table.indexes().get(i));
    }
    for (int i = 0; i < table.foreignKeys().size(); i++) {
      recordBuilder.prop("spannerForeignKey_" + i, table.foreignKeys().get(i));
    }
    SchemaBuilder.FieldAssembler<Schema> fieldsAssembler = recordBuilder.fields();
    for (Column cm : table.columns()) {
      SchemaBuilder.FieldBuilder<Schema> fieldBuilder = fieldsAssembler.name(cm.name());
      fieldBuilder.prop("sqlType", cm.typeString());
      for (int i = 0; i < cm.columnOptions().size(); i++) {
        fieldBuilder.prop("spannerOption_" + i, cm.columnOptions().get(i));
      }
      Schema avroType = avroType(cm.type());
      if (!cm.notNull()) {
        avroType = wrapAsNullable(avroType);
      }
      fieldBuilder.type(avroType).noDefault();
    }
    Schema schema = fieldsAssembler.endRecord();
    schemas.add(schema);
  }

  return schemas;
}
 

private static SchemaBuilder.FieldAssembler<Schema> fieldAvroType(
    final int columnType,
    final int precision,
    final SchemaBuilder.FieldBuilder<Schema> fieldBuilder,
    boolean useLogicalTypes) {

  final SchemaBuilder.BaseTypeBuilder<
          SchemaBuilder.UnionAccumulator<SchemaBuilder.NullDefault<Schema>>>
      field = fieldBuilder.type().unionOf().nullBuilder().endNull().and();

  switch (columnType) {
    case VARCHAR:
    case CHAR:
    case CLOB:
    case LONGNVARCHAR:
    case LONGVARCHAR:
    case NCHAR:
      return field.stringType().endUnion().nullDefault();
    case BIGINT:
      if (precision > 0 && precision <= JdbcAvroRecord.MAX_DIGITS_BIGINT) {
        return field.longType().endUnion().nullDefault();
      } else {
        return field.stringType().endUnion().nullDefault();
      }
    case INTEGER:
    case SMALLINT:
    case TINYINT:
      return field.intType().endUnion().nullDefault();
    case TIMESTAMP:
    case DATE:
    case TIME:
    case TIME_WITH_TIMEZONE:
      if (useLogicalTypes) {
        return field
            .longBuilder()
            .prop("logicalType", "timestamp-millis")
            .endLong()
            .endUnion()
            .nullDefault();
      } else {
        return field.longType().endUnion().nullDefault();
      }
    case BOOLEAN:
      return field.booleanType().endUnion().nullDefault();
    case BIT:
      if (precision <= 1) {
        return field.booleanType().endUnion().nullDefault();
      } else {
        return field.bytesType().endUnion().nullDefault();
      }
    case BINARY:
    case VARBINARY:
    case LONGVARBINARY:
    case ARRAY:
    case BLOB:
      return field.bytesType().endUnion().nullDefault();
    case DOUBLE:
      return field.doubleType().endUnion().nullDefault();
    case FLOAT:
    case REAL:
      return field.floatType().endUnion().nullDefault();
    default:
      return field.stringType().endUnion().nullDefault();
  }
}