Examples with Column com.google.cloud.teleport.spanner.ddl.Column used on opensource projects

Example 1 with Column

use of com.google.cloud.teleport.spanner.ddl.Column in project DataflowTemplates by GoogleCloudPlatform.

the class AvroRecordConverter method apply.

public Mutation apply(GenericRecord record) {
    Schema schema = record.getSchema();
    List<Schema.Field> fields = schema.getFields();
    Mutation.WriteBuilder builder = Mutation.newInsertOrUpdateBuilder(table.name());
    for (Schema.Field field : fields) {
        String fieldName = field.name();
        Column column = table.column(fieldName);
        if (column == null) {
            throw new IllegalArgumentException(String.format("Cannot find corresponding column for field %s in table %s schema %s", fieldName, table.prettyPrint(), schema.toString(true)));
        }
        if (column.isGenerated()) {
            // Spanner will compute generated column values automatically.
            continue;
        }
        Schema avroFieldSchema = field.schema();
        if (avroFieldSchema.getType() == Schema.Type.UNION) {
            Schema unpacked = AvroUtil.unpackNullable(avroFieldSchema);
            if (unpacked != null) {
                avroFieldSchema = unpacked;
            }
        }
        LogicalType logicalType = LogicalTypes.fromSchema(avroFieldSchema);
        Schema.Type avroType = avroFieldSchema.getType();
        switch(column.type().getCode()) {
            case BOOL:
            case PG_BOOL:
                builder.set(column.name()).to(readBool(record, avroType, fieldName).orElse(null));
                break;
            case INT64:
            case PG_INT8:
                builder.set(column.name()).to(readInt64(record, avroType, fieldName).orElse(null));
                break;
            case FLOAT64:
            case PG_FLOAT8:
                builder.set(column.name()).to(readFloat64(record, avroType, fieldName).orElse(null));
                break;
            case STRING:
            case PG_VARCHAR:
            case PG_TEXT:
            case JSON:
                builder.set(column.name()).to(readString(record, avroType, fieldName).orElse(null));
                break;
            case BYTES:
            case PG_BYTEA:
                builder.set(column.name()).to(readBytes(record, avroType, fieldName).orElse(null));
                break;
            case TIMESTAMP:
            case PG_TIMESTAMPTZ:
                builder.set(column.name()).to(readTimestamp(record, avroType, logicalType, fieldName).orElse(null));
                break;
            case DATE:
            case PG_DATE:
                builder.set(column.name()).to(readDate(record, avroType, logicalType, fieldName).orElse(null));
                break;
            case NUMERIC:
                builder.set(column.name()).to(readNumeric(record, avroType, fieldName).orElse(null));
                break;
            case PG_NUMERIC:
                builder.set(column.name()).to(Value.pgNumeric(readPgNumeric(record, avroType, fieldName).orElse(null)));
                break;
            case ARRAY:
            case PG_ARRAY:
                {
                    Schema arraySchema = avroFieldSchema.getElementType();
                    if (arraySchema.getType() == Schema.Type.UNION) {
                        Schema unpacked = AvroUtil.unpackNullable(arraySchema);
                        if (unpacked != null) {
                            arraySchema = unpacked;
                        }
                    }
                    LogicalType arrayLogicalType = LogicalTypes.fromSchema(arraySchema);
                    Schema.Type arrayType = arraySchema.getType();
                    switch(column.type().getArrayElementType().getCode()) {
                        case BOOL:
                        case PG_BOOL:
                            builder.set(column.name()).toBoolArray(readBoolArray(record, arrayType, fieldName).orElse(null));
                            break;
                        case INT64:
                        case PG_INT8:
                            builder.set(column.name()).toInt64Array(readInt64Array(record, arrayType, fieldName).orElse(null));
                            break;
                        case FLOAT64:
                        case PG_FLOAT8:
                            builder.set(column.name()).toFloat64Array(readFloat64Array(record, arrayType, fieldName).orElse(null));
                            break;
                        case STRING:
                        case PG_VARCHAR:
                        case PG_TEXT:
                        case JSON:
                            builder.set(column.name()).toStringArray(readStringArray(record, arrayType, fieldName).orElse(null));
                            break;
                        case BYTES:
                        case PG_BYTEA:
                            builder.set(column.name()).toBytesArray(readBytesArray(record, arrayType, fieldName).orElse(null));
                            break;
                        case TIMESTAMP:
                        case PG_TIMESTAMPTZ:
                            builder.set(column.name()).toTimestampArray(readTimestampArray(record, arrayType, arrayLogicalType, fieldName).orElse(null));
                            break;
                        case DATE:
                        case PG_DATE:
                            builder.set(column.name()).toDateArray(readDateArray(record, arrayType, fieldName).orElse(null));
                            break;
                        case NUMERIC:
                            builder.set(column.name()).toStringArray(readNumericArray(record, arrayType, fieldName).orElse(null));
                            break;
                        case PG_NUMERIC:
                            builder.set(column.name()).toPgNumericArray(readPgNumericArray(record, arrayType, fieldName).orElse(null));
                            break;
                        default:
                            throw new IllegalArgumentException(String.format("Cannot convert field %s in schema %s table %s", fieldName, schema.toString(true), table.prettyPrint()));
                    }
                    break;
                }
            default:
                throw new IllegalArgumentException(String.format("Cannot convert field %s in schema %s table %s", fieldName, schema.toString(true), table.prettyPrint()));
        }
    }
    return builder.build();
}

Example 2 with Column

use of com.google.cloud.teleport.spanner.ddl.Column in project DataflowTemplates by GoogleCloudPlatform.

the class AvroSchemaToDdlConverter method toTable.

public Table toTable(String tableName, Schema schema) {
    if (tableName == null) {
        tableName = schema.getName();
    }
    LOG.debug("Converting to Ddl tableName {}", tableName);
    Table.Builder table = Table.builder(dialect);
    table.name(tableName);
    for (Schema.Field f : schema.getFields()) {
        Column.Builder column = table.column(f.name());
        String sqlType = f.getProp("sqlType");
        String expression = f.getProp("generationExpression");
        if (expression != null) {
            // This is a generated column.
            if (Strings.isNullOrEmpty(sqlType)) {
                throw new IllegalArgumentException("Property sqlType is missing for generated column " + f.name());
            }
            String notNull = f.getProp("notNull");
            if (notNull == null) {
                throw new IllegalArgumentException("Property notNull is missing for generated column " + f.name());
            }
            column.parseType(sqlType).notNull(Boolean.parseBoolean(notNull)).generatedAs(expression);
            String stored = f.getProp("stored");
            if (stored == null) {
                throw new IllegalArgumentException("Property stored is missing for generated column " + f.name());
            }
            if (Boolean.parseBoolean(stored)) {
                column.stored();
            }
        } else {
            boolean nullable = false;
            Schema avroType = f.schema();
            if (avroType.getType() == Schema.Type.UNION) {
                Schema unpacked = unpackNullable(avroType);
                nullable = unpacked != null;
                if (nullable) {
                    avroType = unpacked;
                }
            }
            if (Strings.isNullOrEmpty(sqlType)) {
                Type spannerType = inferType(avroType, true);
                sqlType = toString(spannerType, true);
            }
            String defaultExpression = f.getProp("defaultExpression");
            column.parseType(sqlType).notNull(!nullable).defaultExpression(defaultExpression);
        }
        ImmutableList.Builder<String> columnOptions = ImmutableList.builder();
        for (int i = 0; ; i++) {
            String spannerOption = f.getProp("spannerOption_" + i);
            if (spannerOption == null) {
                break;
            }
            columnOptions.add(spannerOption);
        }
        column.columnOptions(columnOptions.build());
        column.endColumn();
    }
    for (int i = 0; ; i++) {
        String spannerPrimaryKey = schema.getProp("spannerPrimaryKey_" + i);
        if (spannerPrimaryKey == null) {
            break;
        }
        if (spannerPrimaryKey.endsWith(" ASC")) {
            String name = spannerPrimaryKey.substring(0, spannerPrimaryKey.length() - 4);
            table.primaryKey().asc(unescape(name, dialect)).end();
        } else if (spannerPrimaryKey.endsWith(" DESC")) {
            String name = spannerPrimaryKey.substring(0, spannerPrimaryKey.length() - 5);
            table.primaryKey().desc(unescape(name, dialect)).end();
        } else {
            throw new IllegalArgumentException("Cannot parse spannerPrimaryKey " + spannerPrimaryKey);
        }
    }
    table.indexes(getNumberedPropsWithPrefix(schema, "spannerIndex_"));
    table.foreignKeys(getNumberedPropsWithPrefix(schema, "spannerForeignKey_"));
    table.checkConstraints(getNumberedPropsWithPrefix(schema, "spannerCheckConstraint_"));
    // Table parent options.
    String spannerParent = schema.getProp("spannerParent");
    if (!Strings.isNullOrEmpty(spannerParent)) {
        table.interleaveInParent(spannerParent);
        // Process the on delete action.
        String onDeleteAction = schema.getProp("spannerOnDeleteAction");
        if (onDeleteAction == null) {
            // Preserve behavior for old versions of exporter that did not provide the
            // spannerOnDeleteAction property.
            onDeleteAction = "no action";
        }
        if (onDeleteAction.equals("cascade")) {
            table.onDeleteCascade();
        } else if (!onDeleteAction.equals("no action")) {
            // This is an unknown on delete action.
            throw new IllegalArgumentException("Unsupported ON DELETE action " + onDeleteAction);
        }
    }
    return table.build();
}

Example 3 with Column

use of com.google.cloud.teleport.spanner.ddl.Column in project DataflowTemplates by GoogleCloudPlatform.

the class DdlToAvroSchemaConverter method convert.

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.dialect() == Dialect.GOOGLE_STANDARD_SQL) {
            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());
            }
        } else if (table.dialect() == Dialect.POSTGRESQL) {
            if (table.primaryKeys() != null) {
                String encodedPk = table.primaryKeys().stream().map(c -> "\"" + c.name() + "\"").collect(Collectors.joining(", "));
                recordBuilder.prop("spannerPrimaryKey", encodedPk);
            }
            for (int i = 0; i < table.primaryKeys().size(); i++) {
                IndexColumn pk = table.primaryKeys().get(i);
                recordBuilder.prop("spannerPrimaryKey_" + i, "\"" + pk.name() + "\" ASC");
            }
        }
        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));
        }
        for (int i = 0; i < table.checkConstraints().size(); i++) {
            recordBuilder.prop("spannerCheckConstraint_" + i, table.checkConstraints().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));
            }
            if (cm.isGenerated()) {
                fieldBuilder.prop("notNull", Boolean.toString(cm.notNull()));
                fieldBuilder.prop("generationExpression", cm.generationExpression());
                fieldBuilder.prop("stored", Boolean.toString(cm.isStored()));
                // Make the type null to allow us not export the generated column values,
                // which are semantically logical entities.
                fieldBuilder.type(SchemaBuilder.builder().nullType()).withDefault(null);
            } else {
                if (cm.defaultExpression() != null) {
                    fieldBuilder.prop("defaultExpression", cm.defaultExpression());
                }
                Schema avroType = avroType(cm.type());
                if (!cm.notNull()) {
                    avroType = wrapAsNullable(avroType);
                }
                fieldBuilder.type(avroType).noDefault();
            }
        }
        Schema schema = fieldsAssembler.endRecord();
        schemas.add(schema);
    }
    for (View view : ddl.views()) {
        SchemaBuilder.RecordBuilder<Schema> recordBuilder = SchemaBuilder.record(view.name()).namespace(this.namespace);
        recordBuilder.prop("googleFormatVersion", version);
        recordBuilder.prop("googleStorage", "CloudSpanner");
        recordBuilder.prop("spannerViewQuery", view.query());
        if (view.security() != null) {
            recordBuilder.prop("spannerViewSecurity", view.security().toString());
        }
        schemas.add(recordBuilder.fields().endRecord());
    }
    for (ChangeStream changeStream : ddl.changeStreams()) {
        SchemaBuilder.RecordBuilder<Schema> recordBuilder = SchemaBuilder.record(changeStream.name()).namespace(this.namespace);
        recordBuilder.prop("googleFormatVersion", version);
        recordBuilder.prop("googleStorage", "CloudSpanner");
        recordBuilder.prop(AvroUtil.CHANGE_STREAM_FOR_CLAUSE, changeStream.forClause() == null ? "" : changeStream.forClause());
        if (changeStream.options() != null) {
            for (int i = 0; i < changeStream.options().size(); i++) {
                recordBuilder.prop("spannerOption_" + i, changeStream.options().get(i));
            }
        }
        schemas.add(recordBuilder.fields().endRecord());
    }
    return schemas;
}