All Downloads are FREE. Search and download functionalities are using the official Maven repository.

org.apache.flink.streaming.util.serialization.PulsarSerializationSchemaWrapper Maven / Gradle / Ivy

There is a newer version: 1.12.0
Show newest version
/*
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.flink.streaming.util.serialization;

import org.apache.flink.annotation.PublicEvolving;
import org.apache.flink.api.common.serialization.SerializationSchema;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.streaming.connectors.pulsar.SerializableFunction;
import org.apache.flink.streaming.connectors.pulsar.config.RecordSchemaType;
import org.apache.flink.streaming.connectors.pulsar.internal.IncompatibleSchemaException;
import org.apache.flink.streaming.connectors.pulsar.internal.SchemaTranslator;
import org.apache.flink.streaming.connectors.pulsar.internal.SchemaUtils;
import org.apache.flink.table.types.AtomicDataType;
import org.apache.flink.table.types.DataType;

import org.apache.pulsar.client.api.Schema;
import org.apache.pulsar.client.api.TypedMessageBuilder;

import java.util.Optional;

import static org.apache.flink.util.Preconditions.checkArgument;
import static org.apache.flink.util.Preconditions.checkNotNull;

/**
 * A wrapper that warp flink {@link SerializationSchema} to {@link PulsarSerializationSchema}.
 */
public class PulsarSerializationSchemaWrapper implements PulsarSerializationSchema, PulsarContextAware {
    private final SerializationSchema serializationSchema;
    private final RecordSchemaType recordSchemaType;
    private final SchemaInfoWrapper schemaInfoWrapper;
    private final Class clazz;
    private final DataType dataType;
    private final SchemaMode schemaMode;
    private final SerializableFunction topicExtractor;
    private final SerializableFunction keyExtractor;
    private final SerializableFunction> deliverAtExtractor;

    private PulsarSerializationSchemaWrapper(SerializationSchema serializationSchema,
                                             RecordSchemaType recordSchemaType,
                                             Class clazz,
                                             SchemaInfoWrapper schemaInfoWrapper,
                                             DataType dataType,
                                             SchemaMode schemaMode,
                                             SerializableFunction topicExtractor,
                                             SerializableFunction keyExtractor,
                                             SerializableFunction> deliverAtExtractor) {
        this.serializationSchema = serializationSchema;
        this.recordSchemaType = recordSchemaType;
        this.schemaInfoWrapper = schemaInfoWrapper;
        this.clazz = clazz;
        this.dataType = dataType;
        this.schemaMode = checkNotNull(schemaMode);
        this.topicExtractor = topicExtractor;
        this.keyExtractor = keyExtractor;
        this.deliverAtExtractor = deliverAtExtractor;
    }

    @Override
    public TypeInformation getProducedType() {
        return null;
    }

    @Override
    public Optional getTargetTopic(T element) {
        return Optional.ofNullable(topicExtractor.apply(element));
    }

    @Override
    public void serialize(T element, TypedMessageBuilder messageBuilder) {
        if (keyExtractor != null) {
            messageBuilder.keyBytes(keyExtractor.apply(element));
        }
        if (deliverAtExtractor != null) {
            deliverAtExtractor.apply(element).ifPresent(deliverAt -> messageBuilder.deliverAt(deliverAt));
        }
        messageBuilder.value(element);
    }

    @Override
    public void open(SerializationSchema.InitializationContext context) throws Exception {
        this.serializationSchema.open(context);
    }

    @Override
    public Schema getSchema() {
        try {
            switch (schemaMode) {
                case SPECIAL:
                    checkNotNull(schemaInfoWrapper, "The schema cannot be null in SPECIAL mode");
                    return new FlinkSchema<>(schemaInfoWrapper.getSchemaInfo(), serializationSchema, null);
                case ATOMIC:
                    return new FlinkSchema<>(SchemaTranslator.atomicType2PulsarSchema(dataType).getSchemaInfo(),
                            serializationSchema, null);
                case POJO:
                    return new FlinkSchema<>(
                            SchemaUtils.buildSchemaForRecordClazz(clazz, recordSchemaType).getSchemaInfo(),
                            serializationSchema, null);
                case ROW:
                    return new FlinkSchema<>(SchemaUtils.buildRowSchema(dataType, recordSchemaType),
                            serializationSchema, null);
            }
        } catch (IncompatibleSchemaException e) {
            throw new IllegalStateException(e);
        }
        if (schemaInfoWrapper != null) {
            return new FlinkSchema<>(schemaInfoWrapper.getSchemaInfo(), serializationSchema, null);
        }
        try {
            if (dataType instanceof AtomicDataType) {
                return new FlinkSchema<>(SchemaTranslator.atomicType2PulsarSchema(dataType).getSchemaInfo(),
                        serializationSchema, null);
            } else {
                // for pojo type, use avro or json
                checkNotNull(clazz, "for non-atomic type, you must set clazz");
                checkNotNull(recordSchemaType, "for non-atomic type, you must set recordSchemaType");
                return new FlinkSchema<>(SchemaUtils.buildRowSchema(dataType, recordSchemaType),
                        serializationSchema, null);
            }
        } catch (IncompatibleSchemaException e) {
            throw new RuntimeException(e);
        }
    }

    enum SchemaMode {
        ATOMIC,
        POJO,
        SPECIAL,
        ROW
    }

    /**
     * Builder for {@link PulsarSerializationSchemaWrapper}.
     */
    @PublicEvolving
    public static class Builder {
        private final SerializationSchema serializationSchema;
        private RecordSchemaType recordSchemaType;
        private SchemaInfoWrapper schemaInfoWrapper;
        private Class clazz;
        private DataType dataType;
        private SchemaMode mode;
        private SerializableFunction topicExtractor = (T) -> null;
        private SerializableFunction keyExtractor;
        private SerializableFunction> deliverAtExtractor;

        public Builder(SerializationSchema serializationSchema) {
            this.serializationSchema = serializationSchema;
        }

        public PulsarSerializationSchemaWrapper.Builder useSpecialMode(Schema schema) {
            checkArgument(mode == null, "you can only set one schemaMode");
            this.mode = SchemaMode.SPECIAL;
            this.schemaInfoWrapper = new SchemaInfoWrapper(schema.getSchemaInfo());
            return this;
        }

        public PulsarSerializationSchemaWrapper.Builder useAtomicMode(DataType dataType) {
            checkArgument(mode == null, "you can only set one schemaMode");
            this.mode = SchemaMode.ATOMIC;
            checkArgument(dataType instanceof AtomicDataType, "you must set an atomic dataType");
            this.dataType = dataType;
            return this;
        }

        public PulsarSerializationSchemaWrapper.Builder usePojoMode(Class clazz,
                                                                       RecordSchemaType recordSchemaType) {
            checkArgument(mode == null, "you can only set one schemaMode");
            this.mode = SchemaMode.POJO;
            checkArgument(recordSchemaType != RecordSchemaType.ATOMIC,
                    "cant ues RecordSchemaType.ATOMIC to build pojo type schema");
            this.clazz = clazz;
            this.recordSchemaType = recordSchemaType;
            return this;
        }

        public PulsarSerializationSchemaWrapper.Builder useRowMode(DataType dataType,
                                                                      RecordSchemaType recordSchemaType) {
            checkArgument(mode == null, "you can only set one schemaMode");
            this.mode = SchemaMode.ROW;
            this.dataType = dataType;
            this.recordSchemaType = recordSchemaType;
            return this;
        }

        public PulsarSerializationSchemaWrapper.Builder setTopicExtractor(
                SerializableFunction topicExtractor) {
            this.topicExtractor = topicExtractor;
            return this;
        }

        public PulsarSerializationSchemaWrapper.Builder setKeyExtractor(
                SerializableFunction keyExtractor) {
            this.keyExtractor = keyExtractor;
            return this;
        }

        public PulsarSerializationSchemaWrapper.Builder setDeliverAtExtractor(
            SerializableFunction> deliverAtExtractor) {
            this.deliverAtExtractor = deliverAtExtractor;
            return this;
        }

        public PulsarSerializationSchemaWrapper build() {
            checkNotNull(mode, "Must set mode " +
                    "use useSpecialMode or useAtomicMode or usePojoMode or useRowMode");
            return new PulsarSerializationSchemaWrapper<>(
                    serializationSchema,
                    recordSchemaType,
                    clazz,
                    schemaInfoWrapper,
                    dataType,
                    mode,
                    topicExtractor,
                    keyExtractor,
                    deliverAtExtractor);
        }
    }
}




© 2015 - 2024 Weber Informatics LLC | Privacy Policy