com.alibaba.ververica.connectors.odps.sink.StaticPartitionOdpsWriter Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of ververica-connector-odps Show documentation
There is a newer version: 1.17-vvr-8.0.8
Show newest version
/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 com.alibaba.ververica.connectors.odps.sink;

import org.apache.flink.types.Row;
import org.apache.flink.util.Preconditions;

import com.alibaba.ververica.connectors.odps.OdpsConf;
import com.alibaba.ververica.connectors.odps.util.OdpsUtils;
import com.aliyun.odps.Odps;
import com.aliyun.odps.PartitionSpec;
import com.aliyun.odps.Table;
import com.aliyun.odps.data.RecordWriter;
import com.aliyun.odps.tunnel.TableTunnel;
import com.aliyun.odps.tunnel.io.TunnelBufferedWriter;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.concurrent.atomic.AtomicInteger;

/**
 * Odps writer to write records to a specified non-partitioned odps table or a specified partition.
 */
public class StaticPartitionOdpsWriter extends OdpsWriter {

	private static final Logger LOGGER = LoggerFactory.getLogger(StaticPartitionOdpsWriter.class);

	private final String partition;
	private final boolean isPartitioned;

	private transient TableTunnel.UploadSession uploadSession;
	private transient RecordWriter recordWriter;
	private transient AtomicInteger currentCachedCount = new AtomicInteger(0);

	public StaticPartitionOdpsWriter(
			long flushIntervalMs,
			OdpsConf odpsConf,
			String table,
			String partition) {
		super(flushIntervalMs, odpsConf, table);
		this.partition = partition;
		Odps odps = OdpsUtils.initOdps(odpsConf);
		String project = odpsConf.getProject();
		Table tableClient = odps.tables().get(project, table);
		try {
			isPartitioned = tableClient.isPartitioned();
		} catch (Throwable e) {
			LOGGER.error("Fail to check whether table {}.{} is partitioned table.", project, table);
			throw new RuntimeException(e);
		}
		if (isPartitioned) {
			Preconditions.checkArgument(partition != null, "Partition cannot be null for partition table!");
			// create partition if not exists.
			PartitionSpec partSpec = new PartitionSpec(partition);
			try {
				if (!tableClient.hasPartition(partSpec)) {
					LOGGER.info("Creates a new partition {} in table {}.", partSpec, table);
					tableClient.createPartition(partSpec, true);
				}
			} catch (Throwable e) {
				LOGGER.error("Fail to create partition {} in table {}.", partSpec, table);
				throw new RuntimeException(e);
			}
		}
	}

	@Override
	public void open() {
		LOGGER.info("Opening writer ...");
		super.open();
		this.uploadSession = createUploadSession();
		this.recordWriter = createRecordWriter(uploadSession);
		currentCachedCount = new AtomicInteger(0);
		LOGGER.info("Finish to open writer.");
	}

	@Override
	public synchronized void flush(boolean commit) {
		LOGGER.info("Flushing writer {} ...", commit);
		if (this.recordWriter != null) {
			try {
				((TunnelBufferedWriter) recordWriter).flush();
				currentCachedCount.set(0);
			} catch (Throwable e) {
				LOGGER.error("Fail to flush writer.", e);
				throw new RuntimeException(e);
			}
		}
		if (commit) {
			try {
				LOGGER.info("Commit to odps && sessionId {}.", uploadSession.getId());
				uploadSession.commit();
			} catch (Throwable e) {
				LOGGER.error("Fail to commit to odps && sessionId {}.", uploadSession.getId(), e);
				throw new RuntimeException(e);
			}
			// rebuild session and recordWriter
			this.uploadSession = createUploadSession();
			this.recordWriter = createRecordWriter(uploadSession);
		}
		LOGGER.info("Finish to flush writer.");
	}

	@Override
	public void close() {
		LOGGER.info("Closing writer ...");
		super.close();

		synchronized (this) {
			if (recordWriter != null) {
				try {
					recordWriter.close();
				} catch (Throwable e) {
					LOGGER.error("Fail to close writer!", e);
					throw new RuntimeException(e);
				} finally {
					recordWriter = null;
				}
			}

			if (uploadSession != null) {
				try {
					uploadSession.commit();
				} catch (Throwable e) {
					LOGGER.error("Fail to commit odps upload session!", e);
					throw new RuntimeException(e);
				} finally {
					uploadSession = null;
				}
			}
		}
		LOGGER.info("Finish to close writer.");
	}

	@Override
	public synchronized void checkAndFlush(String staticPartKey) {
		if (currentCachedCount.incrementAndGet() > batchCount) {
			flush(false);
		}
	}

	@Override
	public TableTunnel.UploadSession getUploadSession(String staticPartKey) {
		return uploadSession;
	}

	@Override
	public RecordWriter getRecordWriter(String staticPartKey) {
		return recordWriter;
	}

	@Override
	public String getPartitionKey(Row row) {
		return partition;
	}

	private TableTunnel.UploadSession createUploadSession() {
		if (!isPartitioned) {
			try {
				String project = odpsConf.getProject();
				TableTunnel.UploadSession newUploadSession = tableTunnel.createUploadSession(project, table);
				LOGGER.info("Created upload session id: " + newUploadSession.getId());
				return newUploadSession;
			} catch (Throwable e) {
				LOGGER.error("Fail to create uploadSession! ", e);
				throw new RuntimeException("Fail to create uploadSession! ", e);
			}
		} else {
			return createUploadSession(new PartitionSpec(partition));
		}
	}
}